Search papers, labs, and topics across Lattice.
This paper introduces a Sobolev-trained diffusion policy to warm-start trajectory optimization (TO) solvers, leveraging feedback gains from the solver to improve policy learning. By incorporating first-order information into the diffusion policy's loss function, the policy avoids compounding errors that typically arise from small deviations during rollout. Experiments show that this approach significantly reduces TO solving time by 2x to 20x and enables faster inference with fewer diffusion steps, due to the more informed policy initialization.
Sobolev training of diffusion policies slashes trajectory optimization time by up to 20x by preventing compounding errors, even when learning from limited data.
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\times$ to $20 \times$. Incorporating first-order information enables predictions with fewer diffusion steps, reducing inference latency.
