Diffusion-Based Approximate MPC:
Fast and Consistent Imitation of Multi-Modal Action Distributions

Pau Marquez Julbe1, Julian Nubert1,2, Henrik Hose3, Sebastian Trimpe3, Katherine Kuchenbecker1
Max Planck Institute for Intelligent Systems1, ETH Zürich2, RWTH Aachen University3
Paper Code (Coming Soon)

Abstract

Approximating model predictive control (MPC) using imitation learning (IL) allows for fast control without solving expensive optimization problems online. However, methods that use neural networks in a simple L2-regression setup fail to approximate multi-modal (set-valued) solution distributions caused by local optima found by the numerical solver or non-convex constraints, such as obstacles, significantly limiting the applicability of approximate MPC in practice. We solve this issue by using diffusion models to accurately represent the complete solution distribution (i.e., all modes) at high control rates (more than 1000 Hz). This work shows that diffusion-based AMPC significantly outperforms L2-regression-based approximate MPC for multi-modal action distributions. In contrast to most earlier work on IL, we also focus on running the diffusion-based controller at a higher rate and in joint space instead of end-effector space. Additionally, we propose the use of gradient guidance during the denoising process to consistently pick the same mode in closed loop to prevent switching between solutions. We propose using the cost and constraint satisfaction of the original MPC problem during parallel sampling of solutions from the diffusion model to pick a better mode online. We evaluate our method on the fast and accurate control of a 7-DoF robot manipulator both in simulation and on hardware deployed at 250 Hz, achieving a speedup of more than 70 times compared to solving the MPC problem online and also outperforming the numerical optimization (used for training) in success ratio.

Method Overview


We propose fast sampling of model predictive control (MPC) policies through diffusion-based imitation learning. While classical least-squares methods (LSM) fail to model multi-modal action distributions, diffusion models capture the different modes, allowing the approximation of the multiple local minima generated by the MPC controller. To enable policy deployment, we incorporate gradient guidance to maintain mode consistency throughout closed-loop execution. We also reduce noise injection, resulting in smoother, low-jerk control commands. Moreover, we provide different strategies to select the diffusion model's mode online at high control rates, improving feasibility and MPC cost. We provide simulation and real-world experiments on a 7DOF KUKA LBR4+.

We visualize open-loop end-effector trajectory distributions in the presence of a spherical obstacle (red circle). Unlike least-squares methods (LSM), which struggle with non-convex distributions, diffusion models effectively capture multiple modes, significantly improving non-convex constraint satisfaction.

BibTeX


@misc{julbe2025diffusionbasedapproximatempcfast,
      title={Diffusion-Based Approximate MPC: Fast and Consistent Imitation of Multi-Modal Action Distributions}, 
      author={Pau Marquez Julbe, Julian Nubert, Henrik Hose, Sebastian Trimpe, and Katherine J. Kuchenbecker},
      year={2025},
      eprint={2504.04603},
      archivePrefix={arXiv},
      primaryClass={cs.RO},
      url={https://arxiv.org/abs/2504.04603}, 
}