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

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. Code and videos are available at https://dream-mpc.github.io.

Humans can anticipate what will happen if they act in a certain way, because they develop an internal model of how the world behaves. The goal of world models is to give machines the same ability: a learned model predicts the likely outcomes of actions, so the system can plan without having to try many options in the real world, where trials can be slow, costly, or unsafe. Many current planning methods try hundreds or thousands of possible action sequences, which becomes especially expensive for robots. We introduce Dream-MPC, which starts from a good initial guess, evaluates only few candidate action sequences, and then directly adjusts them to achieve better results. Dream-MPC also accounts for uncertainty in its predictions and reuses information across planning steps to avoid starting over each time. Experiments on 24 simulated robotic tasks show that Dream-MPC improves performance and can outperform state-of-the-art baselines, including gradient-free planning methods.