Diffusion Models for Generating Ballistic Spacecraft Trajectories

Song and Ermon’s key insight is that a neural network can be trained to estimate the score function directly from data using denoising score matching.[24] The immediate question that follows this notion is how to train a model that can accurately estimate the score function from a dataset where we do not know the score a priori. From a high-level perspective, it is necessary that the optimal model would minimize the objective:

where $||\cdot||^{2}_{2}$ is used to represent the squared-$\mathcal{L}_{2}$ norm of a vector. Equation (2) resembles the Fisher divergence[25], which compares the squared $\mathcal{L}_{2}$ distance between the data’s true score and the model predicted score. While intuitive, this objective is non-useful since we do not know the true score. To get around this problem, a family of methods called score matching[26] has been developed to minimize Equation (2) without needing to estimate $p_{\mathrm{data}}(\bm{x})$ first. Following the procedures in previous works and again laid out by Song and Ermon[9], Equation (2) can be rewritten as:

where $\operatorname{tr}(\cdot)$ is the matrix trace operator. Although Equation (3) is only dependent on the neural network $s_{\bm{\theta}}$, it is not scalable to high dimensional data since the operation within the trace is equivalent to computing the Jacobian of $s_{\bm{\theta}}$ which in practice is too large to compute efficiently. To completely circumvent $\operatorname{tr}(\nabla_{\bm{x}}s_{\bm{\theta}}(\bm{x}))$ Song and Ermon[9] suggest adopting denoising score matching[24]. Denoinsing score matching starts by perturbing $p_{\mathrm{data}}(\bm{x})$ by a sequence of noise scales $\{\sigma_{i}\}_{i=1}^{L}$ where $\sigma_{1}>\sigma_{2}>...>\sigma_{L}$ and $\sigma^{2}$ is the variance of the Gaussian noise used to perturb $\bm{x}$. Now, we define the Gaussian noise as $p_{\sigma}(\tilde{\bm{x}}|\bm{x})=\mathcal{N}(\tilde{\bm{x}}|\bm{x},\sigma^{2}\mathbf{I})$, where $\mathcal{N}$ represents the Gaussian Normal Distribution, and can produce the entire noise-perturbed distribution as:

Using this information Song and Ermon[9] propose estimating the score at each noise level $p_{\sigma_{i}}(\bm{x})$ by training the neural net conditioned on $\bm{x}$ and $\sigma$, noted as $s_{\bm{\theta}}(\bm{x},\sigma)$. This updated definition is now called the Noise Conditional Score Network (NCSN) and has an updated loss function:

Using Equation (5), all expectations can now be estimated using empirical averages based on the Monte Carlo (MC) approximation of the expectation. An important property of Equation (5) and the denoising score matching techniques employed here is that no specific neural network architecture is required. The only constraint the loss function places is that the model input and output dimensionality match, which can be satisfied with many different model types. This architectural flexibility is attractive because it allows the designer to design models around hardware constraints directly.

Sampling from the model is done through Annealed Langevin dynamics, which can produce samples from a probability density $p_{\mathrm{data}}(\bm{x})$ using only the trained score function. The implementation of Langevin dynamics for generating samples is shown in Algorithm 1. Three parameters need to be precomputed according to the techniques proposed by Song and Ermon[19]. They are the fixed step size $\epsilon$, sampling steps $T$, a sample $(\bm{x}_{0})$ from any prior distribution $\pi(\bm{x})$ and noise scales $\{\sigma_{i}\}_{i=1}^{L}$. The process is presented in Algorithm 1. and can be understood more simply as noisy gradient ascent[19] on $\log p_{\mathrm{data}}(\bm{x})$.

When reviewing Algorithm 1, the approach can be thought of more intuitively as repeatedly sampling from the noisy distributions $p_{\sigma_{1}}(\bm{x}),p_{\sigma_{2}}(\bm{x}),...,p_{\sigma_{L}}(\bm{x})$. We also define the step schedule as $\alpha_{i}=\epsilon\sigma_{i}^{2}/\sigma_{L}^{2}$ for the $i^{th}$ noise scale. A final denoising step, step 11 in Algorithm 1, then denoises samples from $p_{\sigma_{1}}(\bm{\bm{x}})$. The samples from a given noise scale are then used to initialize the next iteration of Langevin dynamics for the next noise scale.

The first evaluation metric directly compares the model-generated output to its real Lambert solution counterpart. The goal is to use the model output’s initial state, final state, and generated $\mathrm{TOF}$ as the inputs to the Pykep Lambert solver. Using the solver’s outputs, $[\bm{v}_{i},\bm{v}_{f}]$, we can then directly compare the model’s generated velocities at the initial and final states to determine how well the model learned to generate a trajectory that solves Lambert’s problem. This process is outlined below and is repeated for the initial and final states.

1. 1.

   Extract $[\mathrm{TOF},\bm{r}_{i},\bm{r}_{f}]$ from the model output states

2. 2.

   Use the Pykep Lambert’s solver[13] to compute a solution and retrieve $[\bm{v}_{i},\bm{v}_{f}]$

3. 3.

   Compute the velocity magnitude difference between the model predicted velocity and the Lambert solution using $\Delta V=||\bm{v}_{\mathrm{model}}-\bm{v}_{\mathrm{Lambert}}||$, at both times $t_{i}$ and $t_{f}$

While the Lambert solver comparison provides insight into the model’s ability to solve Lambert’s problem, it does not provide useful information about the model’s ability to predict the entire trajectory while obeying the underlying two-body dynamics. To quantify the model’s ability to plan and predict the entire trajectory, we propose a new metric: the Defect RMS Number (DRN), where RMS is the root-mean-square. The key concepts are derived from multiple-shooting optimization routines, which use the defect between nodes in a trajectory as the measure of convergence. Multiple shooting breaks a trajectory into $N$ nodes, where each node has a state, velocity, and, in specific cases, control. We refer the reader to the work done by Ré[34] for further discussion on multiple-shooting implementations. For this work the DRN is computed by finding the defect along a trajectory at the midpoint between two adjacent nodes. The $N-1$ defects at all the midpoints are then scaled by $1\mathrm{AU}$ in position and $30\mathrm{km/s}$ for velocity based on previous analysis by Ré[34]. This is based on requirements for optimization routines to have all variables scaled on the order of 1. Multiple shooting routines are typically considered to be converged when a constraint tolerance of $10^{-10}$ is met[34]. We design the DRN to be a single number representing how close a generated trajectory is to being considered converged. The process for computing the DRN is shown below:

1. 1.

   Compute all midpoint propagation times, $t_{\mathrm{mid}}=\frac{t_{i+1}-t_{i}}{2}$

2. 2.

   Propagate node $i$ forward to time $t_{mid}$, to get state $\bm{X}_{forward}$

3. 3.

   Propagate node $i+1$ backward to time$t_{mid}$, to get state $\bm{X}_{backward}$

4. 4.

   Compute the defect $\bm{D}$ at node $i$, using the state difference, $\bm{D}=\bm{X}_{forward}-\bm{X}_{backward}$

5. 5.

   Scale the defects vector positions by $1\mathrm{AU}$ and velocities by $30\mathrm{km/s}$

6. 6.

   Compute the $\mathrm{DRN}$ of $\bm{D}$ using the root-mean-square, $\mathrm{DRN}=\sqrt{\frac{1}{N-1}\sum_{i=1}^{N-1}\sum_{j=1}^{4}{\bm{D}^{2}_{ij}}}$, where $N$ is the number of nodes in a trajectory

Figure 3 visually displays the forward/backward propagation process for a 16-node generated trajectory. During model evaluation, we compute the DRN for all 1000 generated samples and record the mean and standard deviation. In addition to trajectory feasibility, the DRN can be used to compare how feasibility changes with time step resolution. As mentioned in the dataset generation section, we train NCSNv2 on varying time resolution trajectories: 16, 64, 256, and 1024 time steps. The DRN allows us to quantify the trade between the number of nodes in a trajectory and the feasibility of the generated trajectories.

The Equivalent DRN (EDRN) extends the Defect RMS Number metric to attempt to produce an equivalent performance comparison between trajectories of different resolutions. As mentioned, the DRN gives the user a single number that determines the measure of convergence, or more simply feasibility, of a given trajectory. When discussing different temporal resolutions, comparing how feasible a trajectory is and if the model is learning to represent the underlying dynamics when the temporal resolution increases more accurately is useful. For trajectories with high temporal resolution, say 1024 time steps, the propagation time used to compute the defect is significantly smaller. These shorter propagation times result in smaller defects. Although there are more defect computations in higher resolution trajectories, it is not apparent that using the DRN alone provides insight into the model’s ability to learn to predict the underlying dynamics more accurately.

To this end, the EDRN metric provides direct insight into the relationship between temporal resolution and NCSNv2’s ability to model the two-body dynamics. EDRN is computed to generate an equivalent comparison to the diffusion model’s sixteen time step resolution version. For all temporal resolutions, computing the EDRN is done by collecting the initial and final nodes in the generated sample and then selecting fourteen time steps evenly spaced temporally between the two terminal states. With the samples reduced to sixteen time step resolution we then use the procedure outlined above to compute the DRN score for all 1000 samples to produce the EDRN. The process of condensing a sixty-four-resolution trajectory into a sixteen-resolution trajectory is shown in Figure 4.

This analysis demonstrates a diffusion model’s ability to generate new ballistic trajectories from a learned distribution. Our prototype problem tasked the model with generating ballistic trajectories from Earth to Mars. We evaluate the feasibility of these generated trajectories based on their difference in velocity magnitude at both initial and final states with the equivalent Lambert solution. In addition, we use the midpoint defect between nodes to measure how close to \sayreal the entire generated trajectory is. The generated trajectories’ initial and final velocity error averages approximately 100 m/s from the velocities produced for the same initial and final positions, and $\mathrm{TOF}$, using a Lambert solver. The developed DRN metric evaluates the midpoint defects and is average on the order of $10^{-3}$ for all model variations. While this error level is too high to be considered converged, the trajectories generated by the diffusion model would make viable initial guesses for optimization routines like multiple shooting. These routines are normally considered to be converged once all constraints have been met to within a tolerance of $10^{-10}$. The simplest path forward for incorporating diffusion models into trajectory design systems is to use them directly as automated initial guess tools. A future conditional model could produce many viable initial guesses for the same user-defined inputs in parallel. Since sampling time depends only on model size and time step resolution, hundreds of initial guesses could be produced in seconds, and the user could select specific generated samples with desirable characteristics. This allows the mission designer to trade objectives like time of flight or fuel mass in near real-time.

An important property of the diffusion model is a consistent error across time steps. When computing the DRN for a batch of 1000 trajectories, we can record the defects in position and velocity at each node and average them over the entire batch to show that the defects remain nearly constant in all state elements over the entire trajectory. For the base model output 64 node trajectory, the average error and standard deviation at each node are plotted in Figure 8. From Figure 8, we see that the average error across all 1000 generated trajectories remains nearly constant regardless of the propagation time as represented by the node. This appealing property of diffusion models is directly related to the generation process. The model predicts all states in a trajectory simultaneously rather than sequentially. Since the model is tasked with estimating the score function, its objective, in the simplest sense, is to make the trajectory increasingly similar to a sample from the underlying data as noise is removed. The score is a vector with the same size as the trajectory and is used to adjust all states simultaneously. This results in an even distribution of error across the trajectory without bias toward any specific time along the trajectory.

The only requirement for training set by the model architecture is that the input and output dimensions must be the same shape. Adding additional input rows or channels to the image vector is straightforward to include 3-dimensional data or control inputs. To generate specific trajectories based on user inputs, much work has already been done on conditional generation with diffusion models[17]. The common theme in these works is the addition of a context vector, which in this case would be the problem parameters such as initial and final conditions. The context vector combined with guidance techniques such as cross-attention or conditional instance normalization enables diffusion models to generate conditional samples from the training data distribution. Another aspect of the presented diffusion model is that we include time as one of the model outputs. While we choose constant time steps between nodes, the user may use adaptive steps to incorporate maneuvers such as gravity assists. Finally, multiple diffusion models could be used in sequence, each using the output from the model before it as guidance to increase fidelity or add specific features to the trajectory, like impulsive maneuvers.