PAD-TRO: Projection-Augmented Diffusion for Direct Trajectory Optimization

Recently, diffusion models have gained popularity and attention in trajectory optimization due to their capability of modeling multi-modal probability distributions. However, addressing nonlinear equality constraints, i.e, dynamic feasibility, remains a great challenge in diffusion-based trajectory optimization. Recent diffusion-based trajectory optimization frameworks rely on a single-shooting style approach where the denoised control sequence is applied to forward propagate the dynamical system, which cannot explicitly enforce constraints on the states and frequently leads to sub-optimal solutions. In this work, we propose a novel direct trajectory optimization approach via model-based diffusion, which directly generates a sequence of states. To ensure dynamic feasibility, we propose a gradient-free projection mechanism that is incorporated into the reverse diffusion process. Our results show that, compared to a recent state-of-the-art baseline, our approach leads to zero dynamic feasibility error and approximately 4x higher success rate in a quadrotor waypoint navigation scenario involving dense static obstacles.

💡 Research Summary

PAD‑TRO (Projection‑Augmented Diffusion for Direct Trajectory Optimization) addresses two fundamental shortcomings of existing diffusion‑based trajectory planners: (1) the inability to enforce terminal constraints and dynamic feasibility when using a single‑shooting style that samples only control sequences, and (2) the reliance on soft penalty terms that allow residual dynamic violations, which degrade downstream tracking performance. The authors propose a direct trajectory optimization framework that samples the state trajectory itself during the reverse diffusion process and incorporates a gradient‑free projection step to guarantee that each predicted state lies within the reachable set defined by the system dynamics.

Technically, the method extends the variance‑preserving reverse diffusion used in model‑based diffusion (MBD) by introducing a bi‑level noise schedule σ_{i,t}=s·(1−\barα_i)^{1/2}·δ_t, where i indexes the diffusion step and t indexes the trajectory horizon. The factor δ_t reduces noise for later time steps, allowing more accurate projection of downstream states onto the reachable set of earlier states. At each reverse diffusion step i, a batch of Ns state trajectories \tilde{X}i is drawn from a Gaussian centered at the current estimate \tilde{x}i with covariance defined by σ{i,t}. A weighted mean \bar{x}i is computed using both the trajectory cost p_J(x)=exp(−J(x)/λ) and an exponential collision‑avoidance term g(x)=∑{k}exp{−κ(‖o(x)−o{obs,k}‖^2−r_k^2)}. This weighted mean approximates the score ∇_{\tilde{x}_i}log p_i(\tilde{x}_i) via the standard diffusion update −\tilde{x}_i−√{\barα_i},\bar{x}_i/(1−\barα_i).

The novel projection mechanism is applied after the score‑guided update: for each time step t, the predicted state \tilde{x}t is projected onto an approximate reachable set generated by sampling feasible control inputs and propagating the dynamics f(·). Because the projection is performed by sampling rather than solving a convex optimization problem, it remains gradient‑free and scalable to high‑dimensional nonlinear systems. This step enforces the equality constraint x{t+1}=f(x_t,u_t) exactly, eliminating dynamic feasibility error.

Experiments involve a 6‑DOF quadrotor navigating a densely cluttered 3‑D environment with multiple waypoints. Compared against the recent state‑of‑the‑art model‑based diffusion method DRAX, PAD‑TRO achieves zero terminal error, zero dynamic feasibility violation, and roughly a four‑fold increase in success rate. The bi‑level noise schedule and projection dramatically improve sample efficiency, reducing the number of collision‑prone rollouts and accelerating convergence. Computational time is comparable to DRAX, demonstrating that the added projection step does not incur prohibitive overhead.

In summary, PAD‑TRO contributes three key innovations: (i) direct diffusion of state trajectories rather than control sequences, (ii) a bi‑level noise schedule that balances exploration and convergence across the diffusion and trajectory horizons, and (iii) a gradient‑free projection scheme that enforces nonlinear equality constraints during reverse diffusion. These advances enable high‑quality, dynamically feasible trajectories for complex nonlinear systems and open new avenues for applying diffusion models to constrained optimal control problems in robotics, autonomous driving, and manipulation.