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

arXiv: https://arxiv.org/abs/2504.04603 · HTML: https://mlsanity.com/paper/2504.04603/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 26 of its 41 references that are in the map.
- First posted on arXiv 6 Apr 2025.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Offline RL and imitation learning from data / [Imitation learning methods](https://mlsanity.com/offlinerl/imitation-learning-methods/)

## Builds on (most cited first, 26 of 26)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [On the Continuity of Rotation Representations in Neural Networks](https://mlsanity.com/paper/1812.07035/) (2018)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://mlsanity.com/paper/2305.08891/) (2023)
- [Orbit: A Unified Simulation Framework for Interactive Robot Learning Environments](https://mlsanity.com/paper/2301.04195/) (2023)
- [Efficient Diffusion Training via Min-SNR Weighting Strategy](https://mlsanity.com/paper/2303.09556/) (2023)
- [Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation](https://mlsanity.com/paper/2405.07503/) (2024)
- [Motion Planning Diffusion: Learning and Planning of Robot Motions with Diffusion Models](https://mlsanity.com/paper/2308.01557/) (2023)
- [Feedback in Imitation Learning: The Three Regimes of Covariate Shift](https://mlsanity.com/paper/2102.02872/) (2021)
- [MPC-Net: A First Principles Guided Policy Search](https://mlsanity.com/paper/1909.05197/) (2019)
- [Diffusion Predictive Control with Constraints](https://mlsanity.com/paper/2412.09342/) (2024)
- [Real-time Perception meets Reactive Motion Generation](https://mlsanity.com/paper/1703.03512/) (2017)
- [Deep Learning Alternative to Explicit Model Predictive Control for Unknown Nonlinear Systems](https://mlsanity.com/paper/2011.03699/) (2020)
- [DiffuSolve: Diffusion-based Solver for Non-convex Trajectory Optimization](https://mlsanity.com/paper/2403.05571/) (2024)
- [Safe and Fast Tracking on a Robot Manipulator: Robust MPC and Neural Network Control](https://mlsanity.com/paper/1912.10360/) (2019)
- [Neural Networks for Fast Optimisation in Model Predictive Control: A Review](https://mlsanity.com/paper/2309.02668/) (2023)
- [Approximate non-linear model predictive control with safety-augmented neural networks](https://mlsanity.com/paper/2304.09575/) (2023)
- [Toward Near-Globally Optimal Nonlinear Model Predictive Control via Diffusion Models](https://mlsanity.com/paper/2412.08278/) (2024)
- [Parameter-Adaptive Approximate MPC: Tuning Neural-Network Controllers without Retraining](https://mlsanity.com/paper/2404.05835/) (2024)
- [Deep Learning-Based Model Predictive Control for Resonant Power Converters](https://mlsanity.com/paper/1810.04872/) (2018)
- [The Mini Wheelbot: A Testbed for Learning-based Balancing, Flips, and Articulated Driving](https://mlsanity.com/paper/2502.04582/) (2025)
- [A nonlinear tracking model predictive control scheme for dynamic target signals](https://mlsanity.com/paper/1911.03304/) (2019)

## Built on by (most cited first, 3 of 3)

- [Generative Models From and For Sampling-Based MPC: A Bootstrapped Approach For Adaptive Contact-Rich Manipulation](https://mlsanity.com/paper/2510.14643/) (2025)
- [CAPE: Context-Aware Diffusion Policy Via Proximal Mode Expansion for Collision Avoidance](https://mlsanity.com/paper/2511.22773/) (2025)
- [HumanFlow -- Diffusion-Driven MAV Navigation Among Humans via Tightly-Coupled Motion Tracking, Forecasting, and Control](https://mlsanity.com/paper/2605.25685/) (2026)

## Abstract (arXiv)

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) up to kilohertz sampling rates. 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.