Constraint-Informed Learning for Warm Starting Trajectory Optimization

TL;DR

This work builds upon recent results from decision-focused learning and presents a set of decision-focused loss functions using the notion of merit functions for optimization problems, and shows that training networks with such constraint-informed losses can better encode the structure of the trajectory optimization problem and jointly learn to reconstruct the primal-dual solution while yielding improved constraint satisfaction.

Abstract

Future spacecraft and surface robotic missions require increasingly capable autonomy stacks for exploring challenging and unstructured domains, and trajectory optimization will be a cornerstone of such autonomy stacks. However, the nonlinear optimization solvers required remain too slow for use on relatively resource-constrained flight-grade computers. In this work, we turn towards amortized optimization, a learning-based technique for accelerating optimization run times, and present TOAST: Trajectory Optimization with Merit Function Warm Starts. Offline, using data collected from a simulation, we train a neural network to learn a mapping to the full primal and dual solutions given the problem parameters. Crucially, we build upon recent results from decision-focused learning and present a set of decision-focused loss functions using the notion of merit functions for optimization problems. We show that training networks with such constraint-informed losses can better encode the structure of the trajectory optimization problem and jointly learn to reconstruct the primal-dual solution while yielding improved constraint satisfaction. Through numerical experiments on a Lunar rover problem and a 3-degrees-of-freedom Mars powered descent guidance problem, we demonstrate that TOAST outperforms benchmark approaches in terms of both computation times and network prediction constraint satisfaction.

Figures (12)

• Figure 1: Schematic of the TOAST approach: learning warm starts for optimization using task-relevant merit functions. Warm starts (green) are predicted by the neural network (pink), which is trained offline with a Lagrangian-based merit function.
• Figure 2: Only the control policy and dual variables are predicted using the LSTM, and the state variables are recovered by propagating the controls through the system dynamics. During the learning process, the Lagrangian-based merit function uses the decision variables from the propagated dynamics to evaluate the chosen loss function.
• Figure 3: Mean-squared-error (MSE) for predicting $x = 1$ with constraint $x \leq 2$.
• Figure 4: Sensitivity analysis of loss functions: Lagrangian loss displays stable sensitivity; MSE variations decrease across perturbation scales. Integrated Lagrangian MSE loss combines sensitivities to $u$, $x$, and $\lambda$.
• Figure 5: Optimal trajectory generated by TOAST using an LSTM NN predicted constraint-informed warm start (in grey) to solve the MPC problem for an optimal trajectory (in blue). Green arrows indicate the rover's heading, and obstacles are red circles.

Theorems & Definitions (1)

• Definition 1: First Order Necessary Conditions: KKT Conditions