# Accelerating trajectory optimization with Sobolev-trained diffusion policies

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

## Key facts

- Cited by 1 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 22 of its 39 references that are in the map.
- First posted on arXiv 21 Apr 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [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)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Sobolev Training for Neural Networks](https://mlsanity.com/paper/1706.04859/) (2017)
- [Knowledge Transfer with Jacobian Matching](https://mlsanity.com/paper/1803.00443/) (2018)
- [TaSIL: Taylor Series Imitation Learning](https://mlsanity.com/paper/2205.14812/) (2022)
- [Physics-Driven Data Generation for Contact-Rich Manipulation via Trajectory Optimization](https://mlsanity.com/paper/2502.20382/) (2025)
- [DiffuSolve: Diffusion-based Solver for Non-convex Trajectory Optimization](https://mlsanity.com/paper/2403.05571/) (2024)
- [CACTO: Continuous Actor-Critic with Trajectory Optimization -- Towards global optimality](https://mlsanity.com/paper/2211.06625/) (2022)
- [Automatic Differentiation of Rigid Body Dynamics for Optimal Control and Estimation](https://mlsanity.com/paper/1709.03799/) (2017)
- [Guided Flow Policy: Learning from High-Value Actions in Offline Reinforcement Learning](https://mlsanity.com/paper/2512.03973/) (2025)
- [Highly-Efficient Differentiable Simulation for Robotics](https://mlsanity.com/paper/2409.07107/) (2024)
- [Enforcing the consensus between Trajectory Optimization and Policy Learning for precise robot control](https://mlsanity.com/paper/2209.09006/) (2022)
- [CACTO-SL: Using Sobolev Learning to improve Continuous Actor-Critic with Trajectory Optimization](https://mlsanity.com/paper/2312.10666/) (2023)
- [On the Search for Feedback in Reinforcement Learning](https://mlsanity.com/paper/2002.09478/) (2020)

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

- [VAMPS: Visual and Motor Policies from Sampling-Based Planning](https://mlsanity.com/paper/2610.05331/) (2026)

## Abstract (arXiv)

Trajectory Optimization (TO) solvers exploit known system dynamics to compute locally optimal trajectories through iterative improvements. A downside is that each new problem instance is solved independently; therefore, convergence speed and quality of the solution found depend on the initial trajectory proposed. To improve efficiency, a natural approach is to warm-start TO with initial guesses produced by a learned policy trained on trajectories previously generated by the solver. Diffusion-based policies have recently emerged as expressive imitation learning models, making them promising candidates for this role. Yet, a counterintuitive challenge comes from the local optimality of TO demonstrations: when a policy is rolled out, small non-optimal deviations may push it into situations not represented in the training data, triggering compounding errors over long horizons. In this work, we focus on learning-based warm-starting for gradient-based TO solvers that also provide feedback gains. Exploiting this specificity, we derive a first-order loss for Sobolev learning of diffusion-based policies using both trajectories and feedback gains. Through comprehensive experiments, we demonstrate that the resulting policy avoids compounding errors, and so can learn from very few trajectories to provide initial guesses reducing solving time by 2× to 20 ×. Incorporating first-order information enables predictions with fewer diffusion steps, reducing inference latency.