Gradient-Informed Monte Carlo Fine-Tuning of Diffusion Models for Low-Thrust Trajectory Design

TL;DR

The paper addresses the challenge of identifying Pareto-optimal, low-thrust spacecraft trajectories in the CR3BP by recasting the search as sampling from a distribution over locally optimal costates. It integrates gradient-informed MCMC (MALA and HMC) with a diffusion-model self-supervised training loop, using analytic derivatives from state-transition matrices to compute target gradients efficiently. Empirical results on a Saturn-Titan transfer show that gradient-based samplers dramatically improve feasibility and Pareto-front coverage over prior RWM approaches, with MALA delivering the best trade-off between quality and computational cost. The final diffusion-model fine-tuning step enables scalable generation of high-quality samples and reveals the global structure of the solution space, promising faster, more scalable preliminary mission design.

Abstract

Preliminary mission design of low-thrust spacecraft trajectories in the Circular Restricted Three-Body Problem is a global search characterized by a complex objective landscape and numerous local minima. Formulating the problem as sampling from an unnormalized distribution supported on neighborhoods of locally optimal solutions, provides the opportunity to deploy Markov chain Monte Carlo methods and generative machine learning. In this work, we extend our previous self-supervised diffusion model fine-tuning framework to employ gradient-informed Markov chain Monte Carlo. We compare two algorithms - the Metropolis-Adjusted Langevin Algorithm and Hamiltonian Monte Carlo - both initialized from a distribution learned by a diffusion model. Derivatives of an objective function that balances fuel consumption, time of flight and constraint violations are computed analytically using state transition matrices. We show that incorporating the gradient drift term accelerates mixing and improves convergence of the Markov chain for a multi-revolution transfer in the Saturn-Titan system. Among the evaluated methods, MALA provides the best trade-off between performance and computational cost. Starting from samples generated by a baseline diffusion model trained on a related transfer, MALA explicitly targets Pareto-optimal solutions. Compared to a random walk Metropolis algorithm, it increases the feasibility rate from 17.34% to 63.01% and produces a denser, more diverse coverage of the Pareto front. By fine-tuning a diffusion model on the generated samples and associated reward values with reward-weighted likelihood maximization, we learn the global solution structure of the problem and eliminate the need for a tedious separate data generation phase.

Figures (10)

• Figure 1: Simplified illustration of the sampling framework: starting from a baseline diffusion model with distribution $p_{\tilde{\alpha}}$ that is not aligned with the target distribution $\pi_\alpha$, gradient-based MCMC and supervised fine-tuning are used to train a refined model that closely matches the target distribution.
• Figure 2: Visualization of the forward and reverse diffusion processes for a spacecraft trajectory optimization problem. The distributions depict datasets of control vectors at various stages of the diffusion process with an example trajectory corresponding to the realization of the data-point marked in red.
• Figure 3: $J^k$ and $J^*$ for a randomly selected sample $\boldsymbol{\lambda}^k$. The objective values are plotted along the line defined by $\boldsymbol{\lambda}^k+s\nabla J^k(\boldsymbol{\lambda})|_{\boldsymbol{\lambda}=\boldsymbol{\lambda}^k}$, where $s\in\mathbb{R}$. The plot on the right is a close-up of the left plot around $s=0$.
• Figure 4: Reformulating optimization into sampling: minima of $J^*(\boldsymbol{\lambda})$ correspond to peaks of the unnormalized target density $\pi_{\alpha}(\boldsymbol{\lambda})\equiv\exp(-\beta J^*(\boldsymbol{\lambda}))$. The scaling factor $\beta$ controls the thickness and magnitude of peaks.
• Figure 5: Flowchart of variables and equations for gradient‐based MCMC sampling.