Dream-MPC: Gradient-Based Model Predictive Control with Latent Imagination

Abstract: State-of-the-art model-based Reinforcement Learning (RL) approaches either use gradient-free, population-based methods for planning, learned policy networks, or a combination of policy networks and planning. Hybrid approaches that combine Model Predictive Control (MPC) with a learned model and a policy prior to leverage the advantages of both paradigms have shown promising results. However, these approaches typically rely on gradient-free optimization methods, which can be computationally expensive for high-dimensional control tasks. While gradient-based methods are a promising alternative, recent works have empirically shown that gradient-based methods often perform worse than their gradient-free counterparts. We propose Dream-MPC, a novel approach that generates few candidate trajectories from a rolled-out policy and optimizes each trajectory by gradient ascent using a learned world model, uncertainty regularization and amortization of optimization iterations over time by reusing previously optimized actions. Our results on 24 continuous control tasks show that Dream-MPC can significantly improve the performance of the underlying policy and can outperform gradient-free MPC and state-of-the-art baselines. We will open source our code and more at https://dream-mpc.github.io.

• The paper presents a hybrid method using policy priors and gradient optimization in latent space to significantly reduce computational burden in MPC.
• It integrates uncertainty regularization and partial action reuse, leading to improved sample efficiency and robustness across diverse continuous control benchmarks.
• Empirical results show Dream-MPC outperforming sampling-based approaches, achieving state-of-the-art scores with markedly fewer model evaluations.

Dream-MPC: Gradient-Based Model Predictive Control with Latent Imagination

Introduction and Motivation

Model-based RL (MBRL) algorithms have made significant progress in their ability to achieve high sample efficiency and generalize across diverse environments by leveraging learned world models to generate imaginary rollouts. However, a persistent challenge is the computational burden of model-predictive planning in high-dimensional control settings. The prevalent paradigm—sampling-based population optimizers such as CEM or MPPI—require hundreds to thousands of sequence rollouts per step, making them unsuitable for many real-world or embedded applications with limited compute. Gradient-based trajectory optimization presents a theoretically attractive alternative due to its scaling properties, but has so far underperformed its gradient-free competitors. The core of the issue remains in both the poor gradient landscape and insufficient warm-start strategies.

The Dream-MPC approach addresses these bottlenecks by proposing a principled hybrid method: leveraging a learned policy prior to sample informative candidate trajectories (latently rolled out), then performing gradient-based optimization in the world model's latent space with explicit uncertainty regularization and amortized optimization via partial action reuse. The mechanism bridges the efficiency of local gradient search with the robustness and prior structure encoded by learned policies.

Methodology

Dream-MPC operates within a classical MBRL/MPC framework, wherein at each decision step, the agent seeks to maximize the expected return over a receding planning horizon in a differentiable latent world model.

• Candidate Generation: Instead of numerous random Gaussian sequences, only a handful ($N$) of trajectories are generated via policy rollouts in the latent space, optionally perturbed if the underlying policy is deterministic.
• Gradient Optimization: Each candidate trajectory is iteratively refined by gradient ascent ($I$ steps) on a model-based estimate of expected return, computed in the compressed latent space of the learned world model.
• Action Reuse (Amortization): To improve both sample and computational efficiency, actions optimized in preceding time steps are partially reused, weighted by a reuse coefficient $\rho$.
• Uncertainty Regularization: To prevent model exploitation and reduce error accumulation, an epistemic uncertainty penalty is added to the optimization objective, using ensemble value prediction variance.
• Selection: The highest-scoring refined trajectory yields the next real environment action; the process repeats every environment step.

  Figure 1: Overview of Dream-MPC—candidate action trajectories from the policy prior are rolled out/optimized in latent space, with gradient-based updates, reuse, and uncertainty regularization.

Dream-MPC's architecture and optimization pipeline is directly compatible with latent MBRL algorithms such as TD-MPC2, BMPC, and Dreamer, highlighting its generality across both control-centric and generative model formulations.

Empirical Results

The paper conducts extensive empirical evaluation across 24 continuous control benchmarks, including the DeepMind Control Suite, MetaWorld, and HumanoidBench. This suite spans both low- and high-dimensional state/action spaces, tasks with dense and sparse rewards, and both proprioceptive and image-based observations. Benchmarks include integration with both TD-MPC2 (control-centric, implicit world models) and Dreamer (generative, reconstruction-based world models).

Aggregate Performance

• State-of-the-Art Scores: Dream-MPC combined with a strong policy network (in BMPC) yields the highest interquartile median and mean normalized scores across all baselines and domains, demonstrating improved sample efficiency and final performance.
• Robustness: Performance profiles indicate not only improved means but reduced variance, suggesting greater consistency across seeds and tasks.

  Figure 2: Aggregate metrics—Dream-MPC with a strong policy achieves the best optimality gap, mean, and median scores.

  

  Figure 3: Score distributions on all tasks reveal reduced variance and higher minimum returns for Dream-MPC.

Notable quantitative improvements are observed, e.g., integrating Dream-MPC with BMPC improves the IQM by 26.7% and the mean by 20.5% over BMPC's policy alone.

Performance Across Domains

• BMPC Integration: Dream-MPC robustly outperforms MPPI and policy-only variants in BMPC, offering both improved mean returns and reduced computational costs—crucial for high-dimensional, low-latency control.
• TD-MPC2 Integration: Performance gain over the policy baseline is significant (144.7% IQM), although the final gap to MPPI depends on prior policy quality and action diversity.
• Visual Observations: Outperforms MPPI for visual-input tasks, confirming the approach's compatibility with more complex perceptual pipelines.

Learning Dynamics

• Training-Time Planning: When deployed already during training (not just inference), Dream-MPC matches or surpasses MPPI (TD-MPC2), with only marginal degradation on the most challenging, high-dimensional tasks, indicating that low-diversity (few) candidate trajectories are generally sufficient.

  Figure 4: Learning curves in DeepMind Control Suite showing that Dream-MPC matches sample-efficiency and asymptotic returns of the best baselines.

Sensitivity and Ablations

• Warm Start Importance: Ablation shows that using the policy prior (rather than a simple Gaussian) for initialization is critical in high-dimensional spaces—naive random starts sharply degrade final return.
• Uncertainty Regularization: Penalizing high value ensemble variance further stabilizes optimization, improving performance in both training and inference-time planning.
• Parameter Sweeps: Sensitivity analysis (Figure 5) demonstrates that Dream-MPC is robust to moderate changes in candidate count, optimization steps, and planning horizon, although careful tuning yields marginal gains.

Figure 6: Ablation studies demonstrate the necessity of policy warm starts and action reuse for Dream-MPC's performance.

Computational Efficiency

• Model Evaluation Reduction: Dream-MPC achieves orders-of-magnitude reduction in world model queries per step (15 for Dream-MPC vs. ~9,000 for MPPI in typical settings).
• Inference Times: Benchmarks show near parity with gradient-free planning in low dimensions and substantial improvements over prior gradient-based approaches.

Theoretical Analysis and Gradients

Gradient signal analysis reveals that Dream-MPC mitigates the primary pathologies that previously limited gradient-based planning: variance explosion with horizon, gradient vanishing/explosion, and poor signal-to-noise ratios. Direct comparison to Grad-MPC indicates that Dream-MPC maintains more stable and informative planner gradients over long rollouts, which is attributed to both the policy warm start and uncertainty regularization.

Implications and Prospects

Practical Implications

Dream-MPC brings gradient-based trajectory optimization closer to practical deployment in embedded and robotics settings. Its reduced model query count and improved return make it a strong candidate for on-device MPC, particularly in settings (such as dexterous robotics or real-time autonomous control) where computational latency is critical. Its compatibility with both visual and state-only pipelines, and its independence from the underlying world model architecture, broaden the class of problems to which it is immediately applicable.

Theoretical Implications

Results challenge the widely-held presumption that gradient-based planning is fundamentally inferior to gradient-free methods in MBRL. The demonstrated gains highlight the importance of warm-start priors (policy-guided initialization), trajectory refinement rather than full-resampling, and explicit regularization for robustness. They echo recent findings that hybrid architectures—combining learned priors, real-time optimization, and model uncertainty—can bridge the gap between sample efficiency, generalization, and control performance.

Future Directions

Dream-MPC paves the way for further advances along several axes:

• Multi-task and Meta-Learning Extensions: Extension to task-conditioned or multi-task world models, as noted by the authors, could enable generalization and transfer in lifelong learning settings.
• Adaptive Parameterization: Dynamic adjustment of candidate count, optimization depth, and uncertainty regularization during training/inference could yield further efficiency gains.
• Integration with Offline RL: The uncertainty-aware objective aligns naturally with conservative offline RL objectives for robust planning from datasets.
• Gradient-Quality Diagnostics: Systematic exploration of gradient properties in deep world models, as in Figure 7/15, informs further architectural and training choices.

Conclusion

Dream-MPC advances the frontier of efficient, model-based planning by demonstrating that gradient-based trajectory optimization with effective warm-starts and uncertainty-aware objectives can consistently match or outperform population-based, sampling-intensive methods. Its sample efficiency, generality, and computational viability mark it as a new standard for high-performance control in continuous, high-dimensional domains.