11institutetext: University of California, Irvine 22institutetext: Yale
University

# MedGen3D: A Deep Generative Framework for Paired 3D Image and Mask
Generation

Kun Han 11 Yifeng Xiong 11 Chenyu You 22 Pooya Khosravi 11 Shanlin sun 11
Xiangyi Yan 11 James Duncan 22 Xiaohui Xie 11

###### Abstract

Acquiring and annotating sufficient labeled data is crucial in developing
accurate and robust learning-based models, but obtaining such data can be
challenging in many medical image segmentation tasks. One promising solution
is to synthesize realistic data with ground-truth mask annotations. However,
no prior studies have explored generating complete 3D volumetric images with
masks. In this paper, we present MedGen3D, a deep generative framework that
can generate paired 3D medical images and masks. First, we represent the 3D
medical data as 2D sequences and propose the Multi-Condition Diffusion
Probabilistic Model (MC-DPM) to generate multi-label mask sequences adhering
to anatomical geometry. Then, we use an image sequence generator and semantic
diffusion refiner conditioned on the generated mask sequences to produce
realistic 3D medical images that align with the generated masks. Our proposed
framework guarantees accurate alignment between synthetic images and
segmentation maps. Experiments on 3D thoracic CT and brain MRI datasets show
that our synthetic data is both diverse and faithful to the original data, and
demonstrate the benefits for downstream segmentation tasks. We anticipate that
MedGen3D’s ability to synthesize paired 3D medical images and masks will prove
valuable in training deep learning models for medical imaging tasks.

###### Keywords:

Deep Generative Framework 3D Volumetric Images with Masks Fidelity and
Diversity Segmentation

##  1 Introduction

In medical image analysis, the availability of a substantial quantity of
accurately annotated 3D data is a prerequisite for achieving high performance
in tasks like segmentation and detection [26, 17, 29, 9, 31, 34, 35, 32, 33].
This, in turn, leads to more precise diagnoses and treatment plans. However,
obtaining and annotating such data presents many challenges, including the
complexity of medical images, the requirement for specialized expertise, and
privacy concerns.

Generating realistic synthetic data presents a promising solution to the above
challenges as it eliminates the need for manual annotation and alleviates
privacy risks. However, most prior studies [16, 6, 7, 4, 3, 20, 28, 34, 32,
33] have focused on 2D image synthesis, with only a few generating
corresponding segmentation masks. For instance, [15] uses dual generative
adversarial networks (GAN) [14, 35] to synthesize 2D labeled retina fundus
images, while [12] combines a label generator [25] with an image generator
[24] to generate 2D brain MRI data. More recently, [27] uses WGAN [5] to
generate small 3D patches and corresponding vessel segmentations.

However, there has been no prior research on generating whole 3D volumetric
images with the corresponding segmentation masks. Generating 3D volumetric
images with corresponding segmentation masks faces two major obstacles. First,
directly feeding entire 3D volumes to neural networks is impractical due to
GPU memory constraints, and downsizing the resolution may compromise the
quality of the synthetic data. Second, treating the entire 3D volume as a
single data point during training is suboptimal because of the limited
availability of annotated 3D data. Thus, innovative methods are required to
overcome these challenges and generate high-quality synthetic 3D volumetric
data with corresponding segmentation masks.

We propose MedGen3D, a novel diffusion-based deep generative framework that
generates paired 3D volumetric medical images and multi-label masks. Our
approach treats 3D medical data as sequences of slices and employs an
autoregressive process to sequentially generate 3D masks and images. In the
first stage, a Multi-Condition Diffusion Probabilistic Model (MC-DPM)
generates mask sequences by combining conditional and unconditional generation
processes. Specifically, the MC-DPM generates mask subsequences (i.e., several
consecutive slices) at any position directly from random noise or by
conditioning on existing slices to generate subsequences forward or backward.
Given that medical images have similar anatomical structures, slice indices
serve as additional conditions to aid the mask subsequence generation. In the
second stage, we introduce a conditional image generator with a seq-to-seq
model from [30] and a semantic diffusion refiner. By conditioning on the mask
sequences generated in the first stage, our image generator synthesizes
realistic medical images aligned with masks while preserving spatial
consistency across adjacent slices.

The main contributions of our work are as follows: 1) Our proposed framework
is the _first_ to address the challenge of synthesizing complete 3D volumetric
medical images with their corresponding masks; 2) we introduce a multi-
condition diffusion probabilistic model for generating 3D anatomical masks
with high fidelity and diversity; 3) we leverage the generated masks to
condition an image sequence generator and a semantic diffusion refiner, which
produces realistic medical images that align accurately with the generated
masks; and 4) we present experimental results that demonstrate the fidelity
and diversity of the generated 3D multi-label medical images, highlighting
their potential benefits for downstream segmentation tasks.

##  2 Preliminary

###  2.1 Diffusion Probabilistic Model

A diffusion probabilistic model (DPM) [18] is a parameterized Markov chain of
length T, which is designed to learn the data distribution $p(X)$. DPM builds
the Forward Diffusion Process (FDP) to get the diffused data point $X_{t}$ at
any time step $t$ by $q\left(X_{t}\mid
X_{t-1}\right)=\mathcal{N}\left(X_{t};\sqrt{1-\beta_{t}}X_{t-1},\beta_{t}I\right)$,
with $X_{0}\sim q(X_{0})$ and $p(X_{T})=\mathcal{N}\left(X_{T};0,I\right)$.
Let $\alpha_{t}=1-\beta_{t}$ and
$\bar{\alpha}_{t}=\prod_{s=1}^{t}\left(1-\beta_{s}\right)$, Reverse Diffusion
Process (RDP) is trained to predict the noise added in the FDP by minimizing:

$Loss(\theta)=\mathbb{E}_{X_{0}\sim
q\left(X_{0}\right),\epsilon\sim\mathcal{N}(0,I),t}\left[\left\|\epsilon-\epsilon_{\theta}\left(\sqrt{\bar{\alpha}_{t}}X_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon,t\right)\right\|^{2}\right],\vspace{-0.5em}$
(1)

where $\epsilon_{\theta}$ is predicted noise and $\theta$ is the model
parameters.

###  2.2 Classifier-free Guidance

Samples from conditional diffusion models can be improved with classifier-free
guidance [19] by setting the condition $c$ as $\emptyset$ with probability
$p$. During sampling, the output of the model is extrapolated further in the
direction of $\epsilon_{\theta}\left(X_{t}\mid c\right)$ and away from
$\epsilon_{\theta}\left(X_{t}\mid\emptyset\right)$ as follows:

$\hat{\epsilon}_{\theta}\left(X_{t}\mid
c\right)=\epsilon_{\theta}\left(X_{t}\mid\emptyset\right)+s\cdot\left(\epsilon_{\theta}\left(X_{t}\mid
c\right)-\epsilon_{\theta}\left(X_{t}\mid\emptyset\right)\right),\vspace{-0.8em}$
(2)

where $\emptyset$ represents a null condition and $s\geq 1$ is the guidance
scale.

##  3 Methodology

We propose a sequential process to generate complex 3D volumetric images with
masks, as illustrated in Figure 1. The first stage generates multi-label
segmentation, and the second stage performs conditional medical image
generation. The details will be presented in the following sections.

Figure 1: Overview of the proposed MedGen3D, including a 3D mask generator to
autoregressively generate the mask sequences starting from a random position
$z$, and a conditional image generator to generate 3D images conditioned on
generated masks.

###  3.1 3D Mask Generator

Due to the limited annotated real data and GPU memory constraints, directly
feeding the entire 3D volume to the network is impractical. Instead, we treat
3D medical data as a series of subsequences. To generate an entire mask
sequence, an initial subsequence of $m$ consecutive slices is unconditionally
generated from random noise. Then the subsequence is expanded forward and
backward in an autoregressive manner, conditioned on existing slices.

Inspired by classifier-free guidance in Section 2.2, we propose a general
Multi-Condition Diffusion Probabilistic Model (MC-DPM) to unify all three
conditional generations (unconditional, forward, and backward). As shown in
Fig. 2, MC-DPM is able to generate mask sequences directly from random noise
or conditioning on existing slices.

Furthermore, as 3D medical data typically have similar anatomical structures,
slices with the same relative position roughly correspond to the same
anatomical regions. Therefore, we can utilize the relative position of slices
as conditions to guide the MC-DPM in generating subsequences of the target
region and control the length of generated sequences.

Figure 2: Proposed 3D mask generator. Given target position $z$, MC-DPM is
designed to generate mask subsequences (length of $m$) for specific region,
unconditionally or conditioning on first or last $n$ slices, according to the
pre-defined probability $p^{C}\in\\{{p_{F},p_{B},p_{U}}\\}$. Binary indicators
are assigned to slices to signify the conditional slices. We ignore the binary
indicators in the inference process for clear visualization with red outline
denoting the conditional slices and green outline denoting the generated
slices.

Train: For a given 3D multi-label mask $M\in\mathbb{R}^{D\times H\times W}$,
subsequneces of $m$ consecutive slices are selected as
$\\{M_{z},M_{z+1},\ldots,M_{z+(m-1)}\\}$, with $z$ as the randomly selected
starting indices. For each subsequence, we determine the conditional slices
$X^{C}\in\\{\mathbb{R}^{n\times H\times W},\emptyset\\}$ by selecting either
the first or the last $n$ slices, or no slice, based on a probability
$p^{C}\in\\{{p_{Forward},p_{Backward},p_{Uncondition}}\\}$. The objective of
the MC-DPM is to generate the remaining slices, denoted as
$X^{P}\in\mathbb{R}^{(m-\text{len}(X^{C}))\times H\times W}$.

To incorporate the position condition, we utilize the relative position of the
subsequence $\tilde{z}=z/D$, where $z$ is the index of the subsequence’s
starting slice. Then we embed the position condition and concatenate it with
the time embedding to aid the generation process. We also utilize a binary
indicator for each slice in the subsequence to signify the existence of
conditional slices.

The joint distribution of reverse diffusion process (RDP) with the conditional
slices $X^{C}$ can be written as:

$p_{\theta}(X_{0:T}^{P}|X^{C},\tilde{z})=p(X_{T}^{P})\prod_{t=1}^{T}p_{\theta}(X_{t-1}^{P}\mid
X_{t}^{P},X^{C},\tilde{z}).\vspace{-1em}$ (3)

where $p(X_{T}^{P})=\mathcal{N}\left(X_{T}^{P};0,I\right)$, $\tilde{z}=z/D$
and $p_{\theta}$ is the distribution parameterized by the model.

Overall, the model will be trained by minimizing the following loss function,
with
$X_{t}^{P}=\sqrt{\bar{\alpha}_{t}}X_{0}^{P}+\sqrt{1-\bar{\alpha}_{t}}\epsilon$:

$\operatorname{Loss}(\theta)=\mathbb{E}_{X_{0}\sim
q(X_{0}),\epsilon\sim\mathcal{N}\left(0,I\right),p^{C},z,t}\left[\left\|\epsilon-\epsilon_{\theta}\left(X_{t}^{P},X^{C},z,t\right)\right\|^{2}\right].\vspace{-0.5em}$
(4)

Inference: During inference, MC-DPM first generates a subsequence of $m$
slices from random noise given a random location $z$. The entire mask sequence
can then be generated autoregressively by expanding in both directions,
conditioned on the existing slices, as shown in Figure 2. Please refer to the
Supplementary for a detailed generation process and network structure.

###  3.2 Conditional Image Generator

In the second step, we employ a sequence-to-sequence method to generate
medical images conditioned on masks, as shown in Figure 3.

Image Sequence Generator: In the sequence-to-sequence generation task, new
slice is the combination of the warped previous slice and newly generated
texture, weighted by a continuous mask [30]. We utilize Vid2Vid [30] as our
image sequence generator. We train Vid2Vid with its original loss, which
includes GAN loss on multi-scale images and video discriminators, flow
estimation loss, and feature matching loss.

Figure 3: Image Sequence Generator. Given the generated 3D mask, the initial
image is generated by Vid2Vid model sequentially. To utilize the semantic
diffusion model (SDM) to refine the initial result, we first apply small steps
(10 steps) noise, and then use three SDMs to refine. The final result is the
mean 3D images from 3 different views (Axial, Coronal, and Sagittal), yielding
significant improvements over the initially generated image.

Semantic Diffusion Refiner: Despite the high cross-slice consistency and
spatial continuity achieved by vid2vid, issues such as blocking, blurriness
and suboptimal texture generation persist. Given that diffusion models have
been shown to generate superior images [11], we propose a semantic diffusion
refiner utilizing a diffusion probabilistic model to refine the previously
generated images.

For each of the 3 different views, we train a semantic diffusion model (SDM),
which takes 2D masks and noisy images as inputs to generate images aligned
with input masks. During inference, we only apply small noising steps (10
steps) to the generated images so that the overall anatomical structure and
spatial continuity are preserved. After that, we refine the images using the
pre-trained semantic diffusion model. The final refined 3D images are the mean
results from 3 views. Experimental results show an evident improvement in the
quality of generated images with the help of semantic diffusion refiner.

##  4 Experiments and Results

###  4.1 Datasets and Setups

Datasets: We conducted experiments on the thoracic site using three thoracic
CT datasets and the brain site with two brain MRI datasets. For both
generative models and downstream segmentation tasks, we utilized the following
datasets:

  * •

SegTHOR [22]: 3D thorax CT scans (25 training, 5 validation, 10 testing);

  * •

OASIS [23]: 3D brain MRI T1 scans (40 training, 10 validation, 10 testing);

For the downstream segmentation task only and the transfer learning, we
utilized 10 fine-tuning, 5 validation, and 10 testing scans from each of the
3D thorax CT datasets of StructSeg-Thorax [2] and Public-Thor [9], as well as
the 3D brain MRI T1 dataset from ADNI [1].

Implementation: For thoracic datasets, we crop and pad CT scans to ($96\times
320\times 320$). The annotations of six organs (left lung, right lung, spinal
cord, esophagus, heart, and trachea) are examined by an experienced radiation
oncologist. We also include a body mask to aid in the image generation of body
regions. For brain MRI datasets, we use Freesurfer [13] to get segmentations
of four regions (cortex, subcortical gray matter, white matter, and CSF), and
then crop the volume to ($192\times 160\times 160$). We assign discrete values
to masks of different regions or organs for both thoracic and brain datasets
and then combine them into one 3D volume. When synthesizing mask sequences, we
resize the width and height of the masks to $128\times 128$ and set the length
of the subsequence $m$ to 6. We use official segmentation models provided by
MONAI[8] along with standard data augmentations, including spatial and color
transformations.

Setup: We compare the synthetic image quality with DDPM [18], 3D-$\alpha$-WGAN
[21] and Vid2Vid [30], and utilize four segmentation models with different
training strategies to demonstrate the benefit for the downstream task.

###  4.2 Evaluate the Quality of Synthetic Image.

Synthetic Dataset: To address the limited availability of annotated 3D medical
data, we used only 30 CT scans from SegTHOR (25 for training and 5 for
validation) and 50 MRI scans from OASIS (40 for training and 10 for
validation) to generate 110 3D thoracic CT scans and 110 3D brain MRI scans,
respectively.

| Thoracic CT | Brain MRI
---|---|---
| FID $\downarrow$ | LPIPS $\uparrow$ | FID $\downarrow$ | LPIPS $\uparrow$
DDPM [18] | 35.2 | 0.316 | 34.9 | 0.298
3D-$\alpha$-WGAN [21] | 136.2 | 0.286 | 136.4 | 0.289
Vid2Vid [30] | 47.3 | 0.300 | 48.2 | 0.324
Ours | 39.6 | 0.305 | 40.3 | 0.326
Table 1: Synthetic image quality comparison between baselines and ours.

We compare the fidelity and diversity of our synthetic data with DDPM [18]
(train 3 for different views), 3D-$\alpha$-WGAN [21], and vid2vid [30] by
calculating the mean Frèchet Inception Distance (FID) and Learned Perceptual
Image Patch Similarity (LPIPS) from 3 different views.

According to Table 1, our proposed method has a slightly lower FID score but a
similar LPIPS score compared to DDPM. We speculate that this is because DDPM
is trained on 2D images without explicit anatomical constraints and only
generates 2D images. On the other hand, 3D-$\alpha$-WGAN [18], which uses much
larger 3D training data (146 for thorax and 414 for brain), has significantly
worse FID and LPIPS scores than our method. Moreover, our proposed method
outperforms Vid2Vid, showing the effectiveness of our semantic diffusion
refiner.

Figure 4: Our proposed method produces more anatomically accurate images
compared to 3D-$\alpha$-WGAN and vid2vid, as demonstrated by the clearer organ
boundaries and more realistic textures. Left: Qualitative comparison between
different generative models. Right: Visualization of synthetic 3D brain MRI
slices at different relative positions.

###  4.3 Evaluate the Benefits for Segmentation Task.

We explore the benefits of synthetic data for downstream segmentation tasks by
comparing Sørensen–Dice coefficient (DSC) of 4 segmentation models, including
Unet2D [26], UNet3D [10], UNETR [17], and Swin-UNETR [29]. In Table 2 and 3,
we utilize real training data (from SegTHOR and OASIS) and synthetic data to
train the segmentation models with 5 different strategies, and test on all 3
thoracic CT datasets and 2 brain MRI datasets. In Table 4, we aim to
demonstrate whether the synthetic data can aid transfer learning with limited
real finetuning data from each of the testing datasets (StructSeg-Thorax,
Public-Thor and ADNI) with four training strategies.

| SegTHOR* | StructSeg-Thorax | Public-Thor
---|---|---|---
|  Unet 2D |  Unet 3D | UNETR |  Swin UNETR |  Unet 2D |  Unet 3D | UNETR |  Swin UNETR |  Unet 2D |  Unet 3D | UNETR |  Swin UNETR
E2-1 | 0.817 | 0.873 | 0.867 | 0.878 | 0.722 | 0.793 | 0.789 | 0.810 | 0.822 | 0.837 | 0.836 | 0.847
E2-2 | 0.815 | 0.846 | 0.845 | 0.854 | 0.736 | 0.788 | 0.788 | 0.803 | 0.786 | 0.838 | 0.814 | 0.842
E2-3 | 0.845 | 0.881 | 0.886 | 0.886 | 0.772 | 0.827 | 0.824 | 0.827 | 0.812 | 0.856 | 0.853 | 0.856
E2-4 | 0.855 | 0.887 | 0.894 | 0.899 | 0.775 | 0.833 | 0.825 | 0.833 | 0.824 | 0.861 | 0.852 | 0.867
E2-5 | 0.847 | 0.891 | 0.890 | 0.897 | 0.783 | 0.833 | 0.823 | 0.835 | 0.818 | 0.864 | 0.858 | 0.867
Table 2: Experiment 2: DSC of different thoracic segmentation models. There
are 5 training strategies, namely: E2-1: Training with real SegTHOR training
data; E2-2: Training with synthetic data; E2-3: Training with both synthetic
and real data; E2-4: Finetuning model from E2-2 using real training data; and
E2-5: finetuning model from E2-3 using real training data. (* denotes the
training data source.)

According to Table 2 and Table 3, the significant DSC difference between 2D
and 3D segmentation models underlines the crucial role of 3D annotated data.
While purely synthetic data (E2-2) fails to achieve the same performance as
real training data (E2-1), the combination of real and synthetic data (E2-3)
improves model performance in most cases, except for Unet2D on the Public-Thor
dataset. Furthermore, fine-tuning the pre-trained model with real data (E2-4
and E2-5) consistently outperforms the model trained only with real data.
Please refer to Supplementary for organ-level DSC comparisons of the Swin-
UNETR model with more details.

| OASIS* | ADNI
---|---|---
|  Unet 2D |  Unet 3D | UNETR |  Swin UNETR |  Unet 2D |  Unet 3D | UNETR |  Swin UNETR
E2-1 | 0.930 | 0.951 | 0.952 | 0.954 | 0.815 | 0.826 | 0.880 | 0.894
E2-2 | 0.905 | 0.936 | 0.935 | 0.934 | 0.759 | 0.825 | 0.828 | 0.854
E2-3 | 0.938 | 0.953 | 0.953 | 0.955 | 0.818 | 0.888 | 0.898 | 0.906
E2-4 | 0.940 | 0.955 | 0.954 | 0.956 | 0.819 | 0.891 | 0.903 | 0.903
E2-5 | 0.940 | 0.954 | 0.954 | 0.956 | 0.819 | 0.894 | 0.902 | 0.906
Table 3: Experiment 2: DSC of brain segmentation models. Please refer to Table
2 for detailed training strategies. (* denotes the training data source.)

According to Table 4, for transfer learning, utilizing the pre-trained model
(E3-2) leads to better performance compared to training from scratch (E3-1).
Additionally, pretraining the model with synthetic data (E3-3 and E3-4) can
facilitate transfer learning to a new dataset with limited annotated data.

|  Thoracic CT |  Brain MRI
---|---|---
|  StructSeg-Thorax* |  Public-Thor* |  ADNI*
E3-1 |  0.845 |  0.897 |  0.946
E3-2 |  0.865 |  0.901 |  0.948
E3-3 |  0.878 |  0.913 |  0.949
E3-4 |  0.882 |  0.914 |  0.949
Table 4: Experiment 3: DSC of Swin-UNETR finetuned with real dataset. There
are 4 training strategies: E3-1: Training from scratch for each dataset using
limited finetuning data; E3-2 Finetuning the model E2-1 from experiment 2;
E3-3 Finetuning the model E2-4 from experiment 2; and E3-4 Finetuning the
model E2-5 from experiment 2. (* denotes the finetuning data source.)

We have included video demonstrations of the generated 3D volumetric images in
the supplementary material, which offer a more comprehensive representation of
the generated image’s quality.

##  5 Conclusion

This paper introduces MedGen3D, a new framework for synthesizing 3D medical
mask-image pairs. Our experiments demonstrate its potential in realistic data
generation and downstream segmentation tasks with limited annotated data.
Future work includes merging the image sequence generator and semantic
diffusion refiner for end-to-end training and extending the framework to
synthesize 3D medical images across modalities. Overall, we believe that our
work opens up new possibilities for generating 3D high-quality medical images
paired with masks, and look forward to future developments in this field.

## References

  * [1] https://adni.loni.usc.edu/
  * [2] https://structseg2019.grand-challenge.org/dataset/
  * [3] Abbasi-Sureshjani, S., Amirrajab, S., Lorenz, C., Weese, J., Pluim, J., Breeuwer, M.: 4d semantic cardiac magnetic resonance image synthesis on xcat anatomical model. In: Medical Imaging with Deep Learning. pp. 6–18. PMLR (2020)
  * [4] Abhishek, K., Hamarneh, G.: Mask2lesion: Mask-constrained adversarial skin lesion image synthesis. In: Simulation and Synthesis in Medical Imaging: 4th International Workshop, SASHIMI 2019, Held in Conjunction with MICCAI 2019, Shenzhen, China, October 13, 2019, Proceedings. pp. 71–80. Springer (2019)
  * [5] Arjovsky, M., Chintala, S., Bottou, L.: Wasserstein gan. arXiv preprint arXiv: Arxiv-1701.07875 (2017)
  * [6] Baur, C., Albarqouni, S., Navab, N.: Melanogans: high resolution skin lesion synthesis with gans. arXiv preprint arXiv:1804.04338 (2018)
  * [7] Bermudez, C., Plassard, A.J., Davis, L.T., Newton, A.T., Resnick, S.M., Landman, B.A.: Learning implicit brain mri manifolds with deep learning. In: Medical Imaging 2018: Image Processing. vol. 10574, pp. 408–414. SPIE (2018)
  * [8] Cardoso, M.J., Li, W., Brown, R., Ma, N., Kerfoot, E., Wang, Y., Murrey, B., Myronenko, A., Zhao, C., Yang, D., et al.: Monai: An open-source framework for deep learning in healthcare. arXiv preprint arXiv:2211.02701 (2022)
  * [9] Chen, X., Sun, S., Bai, N., Han, K., Liu, Q., Yao, S., Tang, H., Zhang, C., Lu, Z., Huang, Q., et al.: A deep learning-based auto-segmentation system for organs-at-risk on whole-body computed tomography images for radiation therapy. Radiotherapy and Oncology 160, 175–184 (2021)
  * [10] Çiçek, Ö., Abdulkadir, A., Lienkamp, S.S., Brox, T., Ronneberger, O.: 3d u-net: learning dense volumetric segmentation from sparse annotation. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19. pp. 424–432. Springer (2016)
  * [11] Dhariwal, P., Nichol, A.: Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems 34, 8780–8794 (2021)
  * [12] Fernandez, V., Pinaya, W.H.L., Borges, P., Tudosiu, P.D., Graham, M.S., Vercauteren, T., Cardoso, M.J.: Can segmentation models be trained with fully synthetically generated data? In: Simulation and Synthesis in Medical Imaging: 7th International Workshop, SASHIMI 2022, Held in Conjunction with MICCAI 2022, Singapore, September 18, 2022, Proceedings. pp. 79–90. Springer (2022)
  * [13] Fischl, B.: Freesurfer. Neuroimage 62(2), 774–781 (2012)
  * [14] Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., Bengio, Y.: Generative adversarial networks. Communications of the ACM 63(11), 139–144 (2020)
  * [15] Guibas, J.T., Virdi, T.S., Li, P.S.: Synthetic medical images from dual generative adversarial networks. arXiv preprint arXiv:1709.01872 (2017)
  * [16] Han, C., Hayashi, H., Rundo, L., Araki, R., Shimoda, W., Muramatsu, S., Furukawa, Y., Mauri, G., Nakayama, H.: Gan-based synthetic brain mr image generation. In: 2018 IEEE 15th international symposium on biomedical imaging (ISBI 2018). pp. 734–738. IEEE (2018)
  * [17] Hatamizadeh, A., Tang, Y., Nath, V., Yang, D., Myronenko, A., Landman, B., Roth, H.R., Xu, D.: Unetr: Transformers for 3d medical image segmentation. In: Proceedings of the IEEE/CVF winter conference on applications of computer vision. pp. 574–584 (2022)
  * [18] Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33, 6840–6851 (2020)
  * [19] Ho, J., Salimans, T.: Classifier-free diffusion guidance. arXiv preprint arXiv: Arxiv-2207.12598 (2022)
  * [20] Kim, S., Kim, B., Park, H.: Synthesis of brain tumor multicontrast mr images for improved data augmentation. Medical Physics 48(5), 2185–2198 (2021)
  * [21] Kwon, G., Han, C., Kim, D.s.: Generation of 3d brain mri using auto-encoding generative adversarial networks. In: Medical Image Computing and Computer Assisted Intervention–MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13–17, 2019, Proceedings, Part III 22. pp. 118–126. Springer (2019)
  * [22] Lambert, Z., Petitjean, C., Dubray, B., Kuan, S.: Segthor: Segmentation of thoracic organs at risk in ct images. In: 2020 Tenth International Conference on Image Processing Theory, Tools and Applications (IPTA). pp. 1–6. IEEE (2020)
  * [23] Marcus, D.S., Wang, T.H., Parker, J., Csernansky, J.G., Morris, J.C., Buckner, R.L.: Open access series of imaging studies (oasis): cross-sectional mri data in young, middle aged, nondemented, and demented older adults. Journal of cognitive neuroscience 19(9), 1498–1507 (2007)
  * [24] Park, T., Liu, M.Y., Wang, T.C., Zhu, J.Y.: Semantic image synthesis with spatially-adaptive normalization. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 2337–2346 (2019)
  * [25] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 10684–10695 (2022)
  * [26] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. pp. 234–241. Springer (2015)
  * [27] Subramaniam, P., Kossen, T., Ritter, K., Hennemuth, A., Hildebrand, K., Hilbert, A., Sobesky, J., Livne, M., Galinovic, I., Khalil, A.A., et al.: Generating 3d tof-mra volumes and segmentation labels using generative adversarial networks. Medical Image Analysis 78, 102396 (2022)
  * [28] Sun, L., Chen, J., Xu, Y., Gong, M., Yu, K., Batmanghelich, K.: Hierarchical amortized gan for 3d high resolution medical image synthesis. IEEE journal of biomedical and health informatics 26(8), 3966–3975 (2022)
  * [29] Tang, Y., Yang, D., Li, W., Roth, H.R., Landman, B., Xu, D., Nath, V., Hatamizadeh, A.: Self-supervised pre-training of swin transformers for 3d medical image analysis. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 20730–20740 (2022)
  * [30] Wang, T.C., Liu, M.Y., Zhu, J.Y., Liu, G., Tao, A., Kautz, J., Catanzaro, B.: Video-to-video synthesis. arXiv preprint arXiv:1808.06601 (2018)
  * [31] Yan, X., Tang, H., Sun, S., Ma, H., Kong, D., Xie, X.: After-unet: Axial fusion transformer unet for medical image segmentation. In: Proceedings of the IEEE/CVF winter conference on applications of computer vision. pp. 3971–3981 (2022)
  * [32] You, C., Dai, W., Liu, F., Su, H., Zhang, X., Staib, L., Duncan, J.S.: Mine your own anatomy: Revisiting medical image segmentation with extremely limited labels. arXiv preprint arXiv:2209.13476 (2022)
  * [33] You, C., Dai, W., Min, Y., Liu, F., Zhang, X., Clifton, D.A., Zhou, S.K., Staib, L.H., Duncan, J.S.: Rethinking semi-supervised medical image segmentation: A variance-reduction perspective. arXiv preprint arXiv:2302.01735 (2023)
  * [34] You, C., Dai, W., Staib, L., Duncan, J.S.: Bootstrapping semi-supervised medical image segmentation with anatomical-aware contrastive distillation. arXiv preprint arXiv:2206.02307 (2022)
  * [35] You, C., Zhao, R., Liu, F., Dong, S., Chinchali, S., Topcu, U., Staib, L., Duncan, J.: Class-aware adversarial transformers for medical image segmentation. In: NeurIPS (2022)


# Revisiting Implicit Differentiation for Learning Problems in Optimal Control

Ming Xu
School of Computing
Australian National University
<EMAIL_ADDRESS>
&Timothy Molloy
School of Engineering
Australian National University
<EMAIL_ADDRESS>
Stephen Gould
School of Computing
Australian National University
<EMAIL_ADDRESS>

###### Abstract

This paper proposes a new method for differentiating through optimal
trajectories arising from non-convex, constrained discrete-time optimal
control (COC) problems using the implicit function theorem (IFT). Previous
works solve a differential Karush-Kuhn-Tucker (KKT) system for the trajectory
derivative, and achieve this efficiently by solving an auxiliary Linear
Quadratic Regulator (LQR) problem. In contrast, we directly evaluate the
matrix equations which arise from applying variable elimination on the
Lagrange multiplier terms in the (differential) KKT system. By appropriately
accounting for the structure of the terms within the resulting equations, we
show that the trajectory derivatives scale linearly with the number of
timesteps. Furthermore, our approach allows for easy parallelization,
significantly improved scalability with model size, direct computation of
vector-Jacobian products and improved numerical stability compared to prior
works. As an additional contribution, we unify prior works, addressing claims
that computing trajectory derivatives using IFT scales quadratically with the
number of timesteps. We evaluate our method on a both synthetic benchmark and
four challenging, learning from demonstration benchmarks including a 6-DoF
maneuvering quadrotor and 6-DoF rocket powered landing.

##  1 Introduction

This paper addresses end-to-end learning problems that arise in the context of
constrained optimal control, including trajectory optimization, inverse
optimal control, and system identification. We propose a novel,
computationally efficient approach to computing analytical derivatives of
state and control trajectories that solve constrained optimal control (COC)
problems with respect to underlying parameters in cost functions, system
dynamics, and constraints (e.g., state and control limits).

The efficient computation of these _trajectory derivatives_ is important in
the (open-loop) solution of optimal control problems for both trajectory
optimization and model predictive control (MPC). Such derivatives are also
crucial in inverse optimal control (also called inverse reinforcement learning
or learning from demonstration), where given expert demonstration
trajectories, the objective is to compute parameters of the cost function that
best explain these trajectories. Extensions of inverse optimal control that
involve inferring parameters of system dynamics in addition to cost functions
also subsume system identification, and provide further motivation for the
efficient computation of trajectory derivatives. Our proposed method of
computing trajectory derivatives enables direct minimization of a loss
function defined over optimal trajectories (e.g., imitation loss) using first-
order descent techniques and is a direct alternative to existing works [2, 20,
22], including those based on bespoke derivations of Pontryagin’s maximum
principle (PMP) in discrete time [20, 22].

Analytical trajectory derivatives for COC problems are derived by
differentiating through the optimality conditions of the underlying
optimization problem and applying Dini’s implicit function theorem.
Derivatives are recovered as the solution to a system of linear equations
commonly referred to as the (differential) KKT system. This framework
underlies all existing works [2, 20, 22] as well as our proposed approach.
Similarly to Amos et al. [2], we construct the differential KKT system,
however we use the identities from Gould et al. [14] that apply variable
elimination on the Lagrange multipliers relating to the dynamics and
constraints. We show how to exploit the block-sparse structure of the
resultant matrix equations to achieve linear complexity in computing the
trajectory derivative with trajectory length. Furthermore, we can parallelize
the computation, yielding superior scalability and numerical stability on
multithreaded systems compared to methods derived from the PMP [20, 22].

Our specific contributions are as follows. First, we derive the analytical
trajectory derivatives for a broad class of constrained (discrete-time)
optimal control problems with additive cost functions and dynamics described
by first-order difference equations. Furthermore, we show that the computation
of these derivatives is linear in trajectory length by exploiting sparsity in
the resulting matrix expressions. Second, we describe how to parallelize the
computation of the trajectory derivatives, yielding lower computation time and
superior numerical stability for long trajectories. Third, we show how to
directly compute vector-Jacobian products (VJPs) with respect to some outer-
level loss over optimal trajectories, yielding further improvements to
computation time in the context of bi-level optimization. This setting
commonly arises in the direct solution of inverse optimal control problems
[17, 33, 34]. Finally, we provide discussion unifying existing methods for
computing trajectory derivatives [2, 20, 22]. We dub our method IDOC (Implicit
Differentiation for Optimal Control) and validate it across numerous
experiments, showing a consistent speedup over existing approaches derived
from the PMP. Furthermore, for constrained problems, IDOC provides
significantly improved trajectory derivatives, resulting in superior
performance for a general learning from demonstration (LfD) task.111Code
available at https://github.com/mingu6/Implicit-Diff-Optimal-Control.

(a) System Figure

(b) Imitation Learning: Rocket Landing

Figure 1: Left a): Our approach addresses learning problems in optimal control
such as imitation learning. IDOC provides a method for computing the
trajectory derivative or alternatively, vector-Jacobian products with respect
to an outer-level loss, important for solving this problem using an end-to-end
learning approach. Right b): IDOC yields superior trajectory derivatives for
the imitation learning task when inequality constraints (rocket tilt and
thrust limits) are present.

##  2 Related Work

#### Differentiating Through Optimal Control Problems.

An optimal control (OC) problem consists of (system) dynamics, a cost
function, an optimal control policy, and a set of state and/or control
constraints. Learning problems in optimal control involve learning some or all
of these aspects, with the problem of learning dynamics referred to as system
identification [29], the problem of learning the cost function referred to as
inverse optimal control (or inverse reinforcement learning or learning from
demonstration) [34, 17, 21, 33, 37, 51], and the problem of learning the
control policy referred to as reinforcement learning [4]. Recent work [20] has
shown that solving these learning problems can be approached in a unified end-
to-end fashion by minimizing task-specific loss functions with respect to
(unknown) parameters of the associated OC problem. This unified formulation
places great importance on the efficient computation of derivatives of optimal
trajectories with respect to said parameters. Methods reliant on computing
trajectory derivatives have, however, been mostly avoided (and argued against)
in the robotics and control literature due to concerns about computational
tractability. For example, bi-level methods of inverse optimal control [34,
33] have mostly been argued against in favor of methods that avoid derivative
calculations by instead seeking to satisfy optimality conditions derived from
KKT [24, 12] or PMP conditions [17, 32, 23, 21].

#### Analytical Trajectory Derivatives.

Our method is a direct alternative to methods such as DiffMPC [2], PDP [20]
and its extension Safe-PDP [22], which differentiate through optimality
conditions to derive trajectory derivatives. Common to these methods is
identifying that trajectory derivatives can be computed by solving an
auxiliary affine-quadratic OC problem, and furthermore, that this can be done
efficiently using a matrix Riccati equation. While our approach still
differentiates the optimality conditions, we avoid solving matrix Riccati
equations and show that this enables easy computation of VJPs, parallelization
across timesteps and improved numerical stability. Like Safe-PDP, IDOC
computes derivatives through COC problems with smooth inequality constraints,
however we show in our experiments that our derivatives are significantly more
stable during training. Section 4.3 provides a detailed description of the
differences between IDOC and existing methods.

#### Differentiable Optimization.

Our work falls under the area of differentiable optimization, which aims to
embed optimization problems within end-to-end learning frameworks. Gould et
al. [14] proposed the deep declarative networks framework, and provide
identities for differentiating through continuous, inequality-constrained
optimization problems. Combinatorial problems [31, 44], as well as specialized
algorithms for convex problems [1] have also been addressed. Differentiable
optimization has been applied to end-to-end learning of state estimation [49,
42] and motion planning problems [6, 2, 26] in robotics, see Pineda et al.
[35] for a recent survey. In addition, numerous machine learning and computer
vision tasks such as pose estimation [9, 40], meta-learning [28], sorting [11,
7] and aligning time series [47] have been investigated under this setting.

##  3 Constrained Optimal Control Formulation

In this section, we provide an overview of the COC problems we consider in
this paper. These problems involve minimizing a cost function with an additive
structure, subject to constraints imposed by system dynamics and additional
arbitrary constraints such as state and control limits. As a result, they can
be formulated as a non-linear program (NLP). We will discuss how to
differentiate through these COC problems in Section 4.

###  3.1 Preliminaries

First, we introduce notation around differentiating vector-valued functions
with respect to vector arguments, consistent with Gould et al. [14]. Let
$f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m}$ be a vector-valued function with
vector arguments and let $\text{D}f\in\mathbb{R}^{n\times m}$ be the (matrix-
valued) derivative where elements are given by

$\left(\text{D}f(x)\right)_{ij}=\frac{\partial f_{i}}{\partial x_{j}}(x).$ (1)

For a scalar-valued function with (multiple) vector-valued inputs
$f:\mathbb{R}^{n}\times\mathbb{R}^{m}\rightarrow\mathbb{R}$ evaluated as
$f(x,y)$, let the second order derivatives
$\text{D}_{XY}^{2}=\text{D}_{X}(\text{D}_{Y}f)^{\top}$. See Gould et al. [14]
for more details.

###  3.2 Optimization Formulation for Constrained Optimal Control

We can formulate the (discrete-time) COC problem as finding a cost-minimizing
trajectory subject to constraints imposed by (possibly non-linear) system
dynamics. In addition, trajectories may be subject to further constraints such
as state and control limits.

To begin, let the dynamics governing the COC system at time $t$ be given by
$x_{t+1}=f_{t}(x_{t},u_{t};\theta)$, where $x_{t}\in\mathbb{R}^{n}$,
$u_{t}\in\mathbb{R}^{m}$ denotes state and control variables,
respectively222The subscript $f_{t}$ allows for time-varying dynamics in
principle, though we do not evaluate IDOC on any optimal control problems that
have this property in our experiments.. For time horizon $T>0$, let
$x\triangleq(x_{0},x_{1},\ldots,x_{T})$ and
$u\triangleq(u_{0},u_{1},\ldots,u_{T-1})$ denote the state and control
trajectories, respectively. Furthermore, let $\theta\in\mathbb{R}^{d}$ be the
parameter vector that parameterizes our COC problem. Our COC problem is then
equivalent to the following constrained optimization problem

$\begin{array}[]{rlr}\text{minimize}&J(x,u;\theta)\triangleq\sum_{t=0}^{T-1}c_{t}(x_{t},u_{t};\theta)+c_{T}(x_{T};\theta)&\\\
\text{subject to}&&\\\ &x_{0}=x_{\text{init}}&\text{(initial state)}\\\
&x_{t+1}-f_{t}(x_{t},u_{t};\theta)=0\quad\forall
t\in\\{0,\ldots,T-1\\}&\text{(dynamics)}\\\ &g_{t}(x_{t},u_{t};\theta)\leq
0\quad\quad\quad\quad\,\forall t\in\\{0,\ldots,T-1\\},&\text{(path ineq.
constraints)}\\\ &h_{t}(x_{t},u_{t};\theta)=0\quad\quad\quad\quad\forall
t\in\\{0,\ldots,T-1\\},&\text{(path eq. constraints)}\\\
&g_{T}(x_{T};\theta)\leq 0&\text{(terminal ineq. constraints)}\\\
&h_{T}(x_{T};\theta)=0,&\text{(terminal eq. constraints)}\end{array}$ (2)

where
$f_{t}:\mathbb{R}^{n}\times\mathbb{R}^{m}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{n}$
describe the dynamics, and
$c_{t}:\mathbb{R}^{n}\times\mathbb{R}^{m}\times\mathbb{R}^{d}\rightarrow\mathbb{R}$,
$c_{T}:\mathbb{R}^{n}\times\mathbb{R}^{d}\ \rightarrow\mathbb{R}$ are the
instantaneous and terminal costs, respectively. Furthermore, the COC system
may be subject to (vector-valued) inequality constraints
$g_{t}:\mathbb{R}^{n}\times\mathbb{R}^{m}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{q_{t}}$
and $g_{T}:\mathbb{R}^{n}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{q_{T}}$
such as control limits and state constraints, for example. Additional equality
constraints
$h_{t}:\mathbb{R}^{n}\times\mathbb{R}^{m}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{s_{t}}$
and $h_{T}:\mathbb{R}^{n}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{s_{T}}$
can also be included.

The decision variables of Equation 2 are the state and control trajectory
$\xi\triangleq(x,u)$, whereas parameters $\theta$ are assumed to be fixed.
Equation 2 can be interpreted as an inequality-constrained optimization
problem with an additive cost structure and vector-valued constraints for the
initial state $x_{0}$, dynamics, etc. over all timesteps. Let
$\xi(\theta)\triangleq(x^{\star},u^{\star})$ be an optimal solution to
Equation 2. We can treat the optimal trajectory $\xi(\theta)$ as an implicit
function of parameters $\theta$, since Equation 2 can be solved to yield an
optimal $\xi(\theta)$ for any (valid) $\theta$.

We can solve Equation 2 for the optimal trajectory $\xi(\theta)$ using a
number of techniques. We can use general purpose solvers [13, 45], as well as
specialized solvers designed to exploit the structure of COC problems, e.g.,
ones based on differential dynamic programming [30, 19]. Regardless, for the
purposes of computing analytical trajectory derivatives using our method (as
well as methods that differentiate through optimality conditions [20, 22, 2]),
we only need to ensure that our solver returns a vector $\xi(\theta)$ which is
a (local) minimizer to Equation 2. We now describe how to differentiate
through these optimal control problems, important in the end-to-end learning
context.

##  4 Trajectory Derivatives using Implicit Differentiation

In this section, we present our identities for computing trajectory
derivatives $\text{D}\xi(\theta)$ based on those derived in Gould et al. [14]
for general optimization problems by leveraging first-order optimality
conditions and the implicit function theorem. We exploit the block structure
of the matrices in these identities that arises in optimal control problems to
enable efficient computation and furthermore, show that computing trajectory
derivatives is linear in the trajectory length $T$.

Before we present the main analytical result, recall the motivation for
computing $\text{D}\xi(\theta)$. Suppose in the LfD context we have a
demonstration trajectory $\xi^{\text{demo}}$ (we can extend this to multiple
trajectories, but choose not to for notational simplicity). Furthermore,
define a loss $\mathcal{L}(\xi(\theta),\xi^{\text{demo}})$ which measures the
deviation from predicted trajectory $\xi(\theta)$ to demonstration trajectory
$\xi^{\text{demo}}$. Ultimately, if we wish to minimize the loss $\mathcal{L}$
with respect to parameters $\theta$ using a first-order decent method, we need
to compute $\text{D}_{\theta}\mathcal{L}(\xi(\theta),\xi^{\text{demo}})$ by
applying the chain rule. Specifically, we compute
$\text{D}_{\theta}\mathcal{L}(\xi)=\text{D}_{\xi}\mathcal{L}(\xi)\text{D}\xi(\theta)$,
which requires trajectory derivative $\text{D}\xi(\theta)$.

###  4.1 Preliminaries

First, we first reorder $\xi$ such that
$\xi=(x_{0},u_{0},x_{1},u_{1},\ldots,x_{T})\in\mathbb{R}^{n_{\xi}\times 1}$,
where $n_{\xi}=(n+m)T+n$. This grouping of decision variable blocks w.r.t.
timesteps is essential for showing linear time complexity of the computation
of the derivative. Let $\xi_{t}\in\mathbb{R}^{n+m}$ represent the subset of
variables in $\xi$ associated to time $t$, with final state
$\xi_{T}=x_{T}\in\mathbb{R}^{n}$. Next, we stack all constraints defined in
Equation 2 into a single vector-valued constraint $r$ comprised of $T+2$
blocks. Specifically, let
$r(\xi;\theta)\triangleq(r_{-1},r_{0},r_{1},\ldots,r_{T})\in\mathbb{R}^{n_{r}}$,
where $n_{r}=(T+1)n+s+q$. The first block is given by $r_{-1}=x_{0}$, and
subsequent blocks are given by

$r_{t}=\begin{cases}(\tilde{g}_{t}(x_{t},u_{t};\theta),h_{t}(x_{t},u_{t};\theta),x_{t+1}-f(x_{t},u_{t};\theta))&\text{for
}0\leq t\leq T-1\\\
(\tilde{g}_{T}(x_{T};\theta),h_{T}(x_{T};\theta))&\text{for }t=T.\end{cases}$
(3)

Here, $\tilde{g}_{t}$ are the subset of active inequality constraints for
$\xi$ (detected numerically using a threshold $\epsilon$) at timestep $t$.
Note, $s=\sum_{t=0}^{T}s_{t}$, $q=\sum_{t=0}^{T}|\tilde{g}_{t}|$ are the total
number of additional equality and active inequality constraints (on top of the
$Tn$ dynamics and $n$ initial state constraints). Each block represents a
group of constraints associated with a particular timestep (except the first
block $r_{-1}$). As we will see shortly, grouping decision variables and
constraints in this way will admit a favorable block-sparse matrix structure
for cost/constraint Jacobians and Hessians required to compute
$\text{D}\xi(\theta)$.

###  4.2 Analytical Results for Trajectory Derivatives

######  Proposition 1 (IDOC).

Consider the optimization problem defined in Equation 2. Suppose $\xi(\theta)$
exists which minimizes Equation 2. Furthermore, assume $f_{t}$, $c_{t}$,
$g_{t}$ and $h_{t}$ are twice differentiable for all $t$ in a neighborhood of
$(\theta,\xi)$. If $\text{rank}(A)=n_{r}$ and furthermore, $H$ is non-
singular, then

$\text{D}\xi(\theta)=H^{-1}A^{\top}(AH^{-1}A^{\top})^{-1}(AH^{-1}B-C)-H^{-1}B,$
(4)

where

$\displaystyle A$
$\displaystyle=\text{D}_{\xi}r(\xi;\theta)\in\mathbb{R}^{n_{r}\times n_{\xi}}$
$\displaystyle B$
$\displaystyle=\text{D}_{\theta\xi}^{2}J(\xi;\theta)-\sum_{t=-1}^{T}\sum_{i=1}^{|r_{t}|}\lambda_{t,i}\text{D}^{2}_{\theta\xi}r_{t}(\xi;\theta)_{i}\in\mathbb{R}^{n_{\xi}\times
d}$ $\displaystyle C$
$\displaystyle=\text{D}_{\theta}r(\xi;\theta)\in\mathbb{R}^{n_{r}\times d}$
$\displaystyle H$
$\displaystyle=\text{D}_{\xi\xi}^{2}J(\xi;\theta)-\sum_{t=-1}^{T}\sum_{i=1}^{|r_{t}|}\lambda_{t,i}\text{D}^{2}_{\xi\xi}r_{t}(\xi;\theta)_{i}\in\mathbb{R}^{n_{\xi}\times
n_{\xi}},$

and Lagrange multipliers $\lambda\in\mathbb{R}^{n_{r}}$ satisfies
$\lambda^{\top}A=\text{D}_{\xi}J(\xi;\theta)$.

###### Proof.

This is a direct application of Proposition 4.5 in Gould et al. [14]. ∎

######  Remark 1.

Matrix $H$ is block diagonal with $T+1$ blocks.

To see this, $\hat{H}=\text{D}_{\xi\xi}^{2}J(\xi;\theta)$ has a block diagonal
structure with blocks given by
$\text{D}_{\xi_{t}\xi_{t}}^{2}c_{t}(\xi_{t};\theta)$, which follows from the
assumption of additive costs in the COC problem described in Equation 2. In
addition, constraints $g_{t}$, $h_{t}$ depend only on $\xi_{t}$ and
furthermore, the dynamics constraint is first order in $\xi_{t+1}$. Therefore,
$\text{D}^{2}_{\xi\xi}r_{t}(\xi;\theta)_{i}$ is only non-zero in the block
relating to $\xi_{t}$ for all $t$ and $i$. We plot the block-sparsity
structure of $H$ in Figure 2(a).

######  Remark 2.

Matrix $A$ is a block-banded matrix that is two blocks wide.

This is shown by noting that the only non-zero blocks in $A$ correspond to
$\text{D}_{\xi_{t}}r_{t}(\xi;\theta)$ and
$\text{D}_{\xi_{t+1}}r_{t}(\xi;\theta)$. The former relates to $g_{t}$,
$h_{t}$ and the $-f_{t}(\xi_{t};\theta)$ component of the dynamics constraint.
The latter only relates to the $x_{t+1}$ component of the dynamics constraint.
Finally, the initial condition block $r_{-1}$ only depends on $\xi_{0}$, hence
the only non-zero block is given by
$\text{D}_{\xi_{0}}r_{-1}(\xi_{0};\theta)$. We plot the block-sparsity
structure of $A$ in Figure 2(b).

######  Proposition 2.

Evaluating Equation 4 has $O(T)$ time complexity.

###### Proof.

From Remarks 1 and 2, $H$ and $A$ have a block diagonal and block-banded
structure with $T+1$ and $2T+2$ blocks, respectively. Therefore, we can
evaluate $H^{-1}A^{\top}$ in $O(T)$ time and furthermore, $H^{-1}A^{\top}$ has
the same structure as $A$. It follows that we can also evaluate $AH^{-1}B-C$
in $O(T)$ time after partitioning $B$ and $C$ into blocks based on groupings
of $\xi$ and $r$.

Next, observe that $AH^{-1}A^{\top}$ yields a block tridiagonal matrix with
$T+2$ blocks on the main diagonal (see Figure 2(c) for a visualization of the
block-sparse structure). As a result, we can evaluate
$(AH^{-1}A^{\top})^{-1}(AH^{-1}B)$ by solving the linear system
$(AH^{-1}A^{\top})Y=AH^{-1}B$ for $Y$ in $O(T)$ time using any linear time
(w.r.t. number of diagonal blocks) block tridiagonal linear solver. A simple
example is the block tridiagonal matrix algorithm, easily derived from block
Gaussian elimination.

Finally, $H^{-1}B$ can easily be evaluated in $O(T)$ time by solving $T+1$
linear systems comprised of the blocks of $H$ and $B$. We conclude that
evaluating Equation 4 requires $O(T)$ operations. ∎

(a) Matrices $H$ and $H^{-1}$

(b) Matrix $A$

(c) Matrix $AH^{-1}A^{\top}$

Figure 2: Block-sparse structure for matrices $H$, $A$, and $AH^{-1}A^{\top}$
assuming $T=3$. For $A$, we combined inequality and equality constraint groups
$g_{t}$ and $h_{t}$ into a single group for simplicity. Shaded regions
indicate (possible) non-zero blocks.

###  4.3 Comparison to DiffMPC and PDP

All existing methods for computing analytical trajectory derivatives involve
solving the linear system

$\underbrace{\begin{bmatrix}H&A^{\top}\\\
A&0\end{bmatrix}}_{K}\begin{bmatrix}\text{D}\xi\\\
-\text{D}\lambda\end{bmatrix}=\begin{bmatrix}-B\\\ -C\end{bmatrix},$ (5)

where blocks $H,A,B,C$ are defined in Section 4.2 and $\text{D}\xi$ is the
desired trajectory derivative. We call Equation 5 the (differential) KKT
system and the matrix $K$ the (differential) KKT matrix. Note, $\text{D}\xi$
is unique if $K$ is non-singular333See Section 10.1 in [8] for further
discussion on conditions for non-singular $K$. To better contextualize IDOC,
we now briefly compare and contrast how previous methods such as DiffMPC [2]
and PDP [20, 22] solve this KKT system.

#### DiffMPC.

DiffMPC [2] proposes a method for differentiating through Linear Quadratic
Regulator (LQR) problems, which are OC problems where the cost function is
(convex) quadratic and the dynamics are affine. The authors show that solving
the LQR problem using matrix Riccati equations can be viewed as an efficient
method for solving a KKT system which encodes the optimality conditions. In
addition, they show that computing trajectory derivatives involves solving a
similar KKT system, motivating efficient computation by solving an auxiliary
LQR problem in the backward pass. The matrix equation interpretation of the
backward pass allows the efficient computation of VJPs for some downstream
loss $\mathcal{L}(\xi)$ in a bi-level optimization context.

DiffMPC extends to handling non-convex OC problems with box constraints on
control inputs via an iterative LQR-based approach. Specificially, such
problems are handled by first computing VJPs w.r.t. the parameters of an LQR
approximation to the non-convex problem around the optimal trajectory
$\xi,\lambda$. Next, the parameters of the approximation are differentiated
w.r.t. the underlying parameters $\theta$ and combined using the chain rule.
However, differentiating the quadratic cost approximation and dynamics
requires evaluating costly higher-order derivatives, e.g.,
$\frac{d}{d\theta}\frac{d^{2}c(\xi)}{d\xi^{2}}$, which are 3D tensors. These
tensors are dense in general, although problem specific sparsity structures
may exist.

#### PDP/Safe-PDP.

PDP [20] and its extension, Safe-PDP [22], take a similar approach to DiffMPC
in deriving trajectory derivatives. PDP derives trajectory derivatives by
starting with PMP, which is well-known in the control community and applies to
non-convex OC problems. Furthermore, Jin et al. [20] shows that the PMP and
KKT conditions are equivalent in the discrete-time COC setting. Trajectory
derivatives are obtained by differentiating the PMP conditions, yielding a new
set of PMP conditions for an auxiliary LQR problem, which can be solved using
matrix Riccati equations. This approach is fundamentally identical to DiffMPC,
with only superficial differences for LQR problems (solving a differential KKT
system versus differentiating PMP conditions)444Jin et al., [20] claim that
the backward pass for DiffMPC is $O(T^{2})$ due to direct inversion of the KKT
matrix. However, DiffMPC actually solves an auxiliary LQR problem for $O(T)$
complexity..

For non-convex problems, PDP does not evaluate higher-order derivatives,
unlike DiffMPC. However, it is not clear how to compute VJPs under approaches
derived using the PDP framework. Safe-PDP extended PDP to handle arbitrary,
smooth inequality constraints using the constrained PMP conditions, which
extends the capability of DiffMPC to handle box constraints. Trajectory
derivatives are computed by solving an auxiliary equality constrained LQR
problem. Experiments in both Jin et al. [22] and Section 6 show that Safe-PDP
yields unreliable derivatives for COC problems. While Jin et al. [22]
conjected that instability was due to the set of active constraints changing
between iterations, we find that IDOC still learns reliably in the presence of
constraint switching. We instead observed that the poor gradient quality
arises naturally from the equality constrained LQR solver used for the
backward pass [25] not matching the solution obtained by directly solving the
KKT system.

####  IDOC.

In contrast, we solve Equation 5 for $\text{D}\xi$ by first applying variable
elimination on $\text{D}\lambda$, yielding Proposition 1. We will discuss in
Section 5 how the equations resulting from this approach admit parallelization
and favorable numerical stability compared to DiffMPC and PDP’s auxiliary LQR
approach. Furthermore, VJPs can be easily computed (unlike PDP) without
evaluating higher-order derivatives (unlike DiffMPC). IDOC handles arbitrary
smooth constraints like Safe-PDP, and we show in Section 6 that our
derivatives for COC problems are reliable during training. However, we require
the additional assumption that $H$ is non-singular (which holds if and only if
all blocks in $H$ are non-singular), which is not required for $K$ to be non-
singular. An oftentimes effective solution when $H$ is singular, initially
proposed by Russell et al. [39], is to set $H=H+\frac{\delta}{2}I$ for small
$\delta$, which is analogous to adding a proximal term to the cost function in
Equation 2. We will now describe how to evaluate Equation 4 in linear time by
exploiting the block-sparse structure of the matrix equation.

##  5 Algorithmic Implications of IDOC

###  5.1 Parallelization and Numerical Stability

#### Parallelization.

To evaluate Equation 4, we can leverage the block diagonal structure of $H$
and block-banded structure of $A$ to compute $H^{-1}A^{\top}$ and $H^{-1}B$ in
parallel across all timesteps. Specifically, we evaluate the matrix product
block-wise, e.g., $H_{t}^{-1}B_{t}$ for all $t$ in parallel (similarly with
$A$). Following this, we can then compute $A(H^{-1}A^{\top})$ in parallel also
using the same argument.

In addition, we can solve the block tridiagonal linear systems involving
$AH^{-1}A^{\top}$ by using a specialized parallel solver [5, 36, 27, 41, 15,
18]. Unfortunately, none of these methods have open source code available, and
so in our experiments, we implement the simple block tridiagonal matrix
algorithm. Implementing a robust parallel solver is an important direction for
future work, and will further improve the scalability and numerical stability
of IDOC.

#### Numerical Stability.

In addition to parallelization, another benefit of avoiding Riccati-style
recursions for computing derivatives is improved numerical stability. For long
trajectories and/or poorly conditioned COC problems (e.g., stiff dynamics),
IDOC reduces the rounding errors that accumulate in recursive approaches. We
show superior numerical stability in our experiments compared to PDP in
Section 6, despite using the simple recursive block tridiagonal matrix
algorithm for solving $AH^{-1}A^{\top}$. Replacing this recursion with a more
sophisticated block tridiagonal solver should further improve the stability of
the backwards pass and is left as future work.

###  5.2 Vector Jacobian Products

Another benefit of explicitly writing out the matrix equations for
$\text{D}\xi(\theta)$ given in Equation 4 is that we can now directly compute
VJPs given some outer loss $\mathcal{L}(\xi)\in\mathbb{R}$ over the optimal
trajectory. Let
$v\triangleq\text{D}_{\xi}\mathcal{L}(\xi)^{\top}\in\mathbb{R}^{n_{\xi}\times
1}$ be the gradient of the loss w.r.t. trajectory $\xi(\theta)$. The desired
gradient $\text{D}_{\theta}\mathcal{L}(\xi(\theta))$ is then given by
$\text{D}_{\theta}\mathcal{L}(\xi)=v^{\top}\text{D}\xi(\theta)$ using the
chain rule. The resultant expression is

$\text{D}_{\theta}\mathcal{L}(\xi(\theta))=v^{\top}(H^{-1}A^{\top}(AH^{-1}A^{\top})^{-1}(AH^{-1}B-C)-H^{-1}B).$
(6)

The simple observation here is that we do not need to construct
$D_{\theta}\xi^{\star}(\theta)$ explicitly, and can instead evaluate the VJP
directly. We propose evaluating Equation 6 from left to right and block-wise,
which will reduce computation time compared to explicitly constructing
$\text{D}\xi(\theta)$ and then multiplying with $v$.

To see this, we follow the example in Gould et al. [14] and assume the blocks
in $H$ have been factored. Then for a single block (ignoring constraints for
simplicity), evaluating $v^{\top}(H^{-1}_{t}B_{t})$ is $O((n+m)^{2}p)$ while
evaluating $(v^{\top}H^{-1}_{t})B_{t}$ is $O((n+m)^{2}+(n+m)p)$. Evaluating
the VJP directly significantly reduces computation time compared to
constructing the full trajectory derivative for problems with higher numbers
of state/control variables and tunable parameters.

##  6 Experiments

While our contributions are largely analytical, we have implemented the
identities in Equation 4 to verify our claims around numerical stability and
computational efficiency on a number of simulated COC environments. We will
show that IDOC is able to compute trajectory derivatives significantly faster
compared to its direct alternative PDP [20] and Safe-PDP [22] with superior
numerical stability.

(a) Cartpole

(b) 6-DoF Quadrotor

(c) Robotarm

(d) 6-DoF Rocket

Figure 3: Learning curves for the imitation learning task over five trials.
Bold lines represent mean loss, lighter lines represent individual trials.
IDOC yields more stable gradients and lower final imitation loss across all
environments and trials compared to both Safe-PDP (S-PDP) and Safe-PDP with
log-barrier functions (S-PDP (b)).

(a) Horizon length scalability.

(b) Model size scalability.

(c) Numerical stability.

Figure 4: Synthetic experiments measuring a) scalability with horizon length
$T$, with fixed parameter size $|\theta|=10$k, b) scalability with number of
parameters $d$ with fixed horizon length $T=1$k, c) numerical stability with
varying block condition number $\kappa(H_{t})$. Error bars measure standard
error across 5 and 25 samples for computation time (a, b) and numerical
stability (c), respectively.

###  6.1 Experimental Setup

We evaluate IDOC against PDP [20] and Safe-PDP [22] in an LfD setting across
four simulation environments proposed in Jin et al. [22], as well as a
synthetic experiment. The simulation environments showcase the gradient
quality for a realistic learning task, whereas the synthetic experiment is
designed to measure numerical stability and computation times. For Safe-PDP,
we evaluate against the method proposed for inequality constrained problems
(Safe-PDP), as well as using an approximate log-barrier problem in place of
the full problem (Safe-PDP (b)). The IPOPT solver [45] is used in the forward
pass to solve the COC problem, and Lagrange multipliers $\lambda$ are
extracted from the solver output. More generally however, note that the method
proposed in Gould et al. [14] can be used to recover $\lambda$ if another
solver is used where $\lambda$ is not provided.

#### Imitation Learning/LfD.

The LfD setting involves recovering the model parameters $\theta$ that
minimizes the mean-squared imitation error to a set of $N$ demonstration
trajectories, given by
$\mathcal{L}(\xi(\theta),\xi^{\text{demo}})\triangleq\frac{1}{N}\sum_{i}\|\xi(\theta)_{i}-\xi^{\text{demo}}_{i}\|^{2}$.
We include all parameters in the cost, dynamics and constraint functions in
$\theta$ and furthermore, evaluate performance both with (S-PDP) and without
(PDP) inequality constraints. We perform experiments in four standard
simulated environments: cartpole, 6-DoF quadrotor555Interestingly, $H$ is
singular in this case, relating to block $H_{T}$, so we use the trick proposed
in Section 4.3 with $\delta=10^{-6}$ to compute Equation 4. More details
around this are provided in the appendix., 2-link robot arm and 6-DoF rocket
landing. For a detailed description of the imitation learning problem as well
as each COC task, see the appendix. In addition, we provide timings for all
methods within the simulation environments in the appendix.

#### Synthetic Benchmark.

The simulation experiments described above are not sufficiently large-scale to
adequately measure the benefits of parallelization and numerical stability
afforded by IDOC over PDP. To demonstrate these benefits, we constructed a
large-scale, synthetic experiment where the blocks required to construct
$H,A,B$ and $C$ (as discussed in Section 4) are generated randomly.
Computation time against the horizon length $T$, as well as the number of
parameters $d$ is reported. Numerical stability is measured by comparing mean-
absolute error between trajectory derivatives evaluated under 32-bit and
64-bit precision for varying condition numbers over the blocks of $H$. State
and control dimensions are fixed at $n=50$ and $m=10$. Experiments are run on
an AMD Ryzen 9 7900X 12-core processor, 128Gb RAM and Ubuntu 20.04. See the
appendix for more details.

###  6.2 Imitation Learning/LfD Results

Figure 5: Control Limit Learning.

In this section, we evaluate the effectiveness of gradients produced from IDOC
against PDP and Safe-PDP [20, 22] for the imitation learning/LfD task. In the
equality constrained setting, IDOC and PDP yield identical (up to machine
precision) trajectory derivatives and imitation loss throughout learning; see
the appendix for more detailed results and analysis. However, as shown in
Figure 3, when inequality constraints are introduced, the derivatives produced
by Safe-PDP and IDOC differ significantly. Safe-PDP fails to reduce the
imitation loss due to unreliable gradients, whereas IDOC successively
decreases the imitation loss. While Safe-PDP (b), which differentiates through
a log-barrier problem, also provides stable learning, it ultimately yields
higher imitation loss compared to IDOC. This is because the barrier problem is
an approximation to the true COC problem. Further discussion around using log-
barrier methods and COC problems is provided in Section 7.1.

###  6.3 Synthetic Benchmark Results

Figure 4 illustrates the results of the synthetic benchmark. We observed that
IDOC and PDP have similar computation time and scalability with problem size
when computing full trajectory derivatives. However, computing VJPs with IDOC
significantly reduces computation time and improves scalability by an order of
magnitude with model size (i.e., $d=|\theta|$). Numerical stability
experiments show that IDOC (full and VJP) yields lower round-off errors and
improved stability compared to PDP. Using more sophisticated block tridiagonal
solvers will further improve stability.

##  7 Discussion and Future Work

###  7.1 Differentiating through Log-Barrier Methods

In the imitation learning setting, we assume that demonstration data
$\xi^{\text{demo}}$ is an optimal solution to a COC problem with (unknown)
parameters $\theta^{\star}$. Furthermore, we assume that a subset of timesteps
$\mathcal{A}\subseteq\\{1,\dots,T\\}$ yield active constraints. Under the log-
barrier formulation, the constraint boundaries must be relaxed beyond their
true values during learning to minimize the imitation loss. Concretely, under
$\theta^{\star}$, $\log
g_{t}(x_{t}^{\text{demo}},u_{t}^{\text{demo}};\theta^{\star})$ are undefined
for $t\in\mathcal{A}$. Therefore to recover $\xi(\theta)=\xi^{\text{demo}}$,
we must have $g_{t}(x_{t}(\theta),u_{t}(\theta))<0$ for $t\in\mathcal{A}$,
i.e., we cannot recover the true constraint function through minimizing the
imitation loss. We verify this using the robot arm environment, and present
the results in Figure 5, plotting the error between the estimated and true
constraint value. We observe that IDOC recovers the true constraint value more
closely compared to Safe-PDP (b).

#### Generality of IFT.

The generality of the differentiable optimization framework, and the matrix
equation formulation for trajectory derivatives given in Equation 4 yield
additional conceptual benefits which may help us tackle even broader classes
of COC problems. For example, we can relax the additive cost assumption and
add a final cost defined over the full trajectory such as

$h_{T}(\xi;\theta)=\xi^{\top}Q\xi,$ (7)

where $Q=CC^{\top}$ and $C\in\mathbb{R}^{n_{\xi}\times k}$ for $k\ll n_{\xi}$
is full rank. For IDOC, we can simply use the matrix inversion lemma to invert
$H$ in $O(T)$. However, this is more complicated for methods derived from the
PMP which rely on the specific additive cost structure of the underlying COC
problem.

#### Dynamic Games.

A generalization of our work is to apply the differentiable optimization
framework to handle learning problems that arise in dynamic games [3] and have
recently begun to employ PDP-based approaches [10]. Being able to
differentiate through solution and equilibrium concepts that arise in dynamic
games enables the solution of problems ranging from minimax robust control [3]
to inverse dynamic games [10, 33], and is a promising direction for future
work.

#### Fast forward passes.

While we have proposed a more efficient way of computing trajectory
derivatives (i.e., the backward pass), we observe that the forward pass to
solve the COC problem bottlenecks the learning process. Significant effort
must be placed on developing fast solvers for COC problems for hardware
accelerators to allow learning to scale to larger problem sizes.

#### Combining MPC and Deep Learning.

A recent application for computing trajectory derivatives is combining ideas
from deep learning (including reinforcement learning) with MPC [38, 50, 46,
48]. Common to these approaches is using a learned model to select the
parameters for the MPC controller. By allowing for robust differentiation
through a broader class COC problems compared to previous approaches, we hope
that IDOC will allow for future development in this avenue of research.

##  8 Conclusion

In this paper, we present IDOC, a novel approach to differentiating through
COC problems. Trajectory derivatives are evaluated by differentiating KKT
conditions and using the implicit function theorem. Contrary to prior works,
we do not solve an auxiliary LQR problem to efficiently solve the
(differential) KKT system for the trajectory derivative. Instead, we apply
variable elimination on the KKT system and solve the resultant matrix
equations directly. We show that linear time evaluation is possible by
appropriately considering equation structure. In fact, we show that IDOC is
faster and more numerically stable in practice compared to methods derived
from the PMP. We hope that our discussion connecting the fields of inverse
optimal control and differentiable optimization will lead to future work which
enables differentiability of a broader class of COC problems.

## References

  * Agrawal et al. [2019] A. Agrawal, B. Amos, S. Barratt, S. Boyd, S. Diamond, and Z. Kolter.  Differentiable convex optimization layers.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Amos et al. [2018] B. Amos, I. Jimenez, J. Sacks, B. Boots, and J. Z. Kolter.  Differentiable MPC for End-to-end Planning and Control.  _Advances in Neural Information Processing Systems_ , 2018.
  * Başar and Olsder [1982] T. Başar and G. J. Olsder.  _Dynamic noncooperative game theory_.  Academic Press, London/New York, 1982.  ISBN 978-0-08-095666-4.
  * Bertsekas [2019] D. Bertsekas.  _Reinforcement learning and optimal control_.  Athena Scientific, 2019.
  * Bevilacqua et al. [1988] R. Bevilacqua, B. Codenotti, and F. Romani.  Parallel solution of block tridiagonal linear systems.  _Linear Algebra and its Applications_ , 104:39–57, 1988\.
  * Bhardwaj et al. [2020] M. Bhardwaj, B. Boots, and M. Mukadam.  Differentiable Gaussian process motion planning.  _Proc. of the IEEE International Conference on Robotics and Automation (ICRA)_ , 2020.
  * Blondel et al. [2020] M. Blondel, O. Teboul, Q. Berthet, and J. Djolonga.  Fast differentiable sorting and ranking.  In _Proc. of the International Conference on Machine Learning (ICML)_ , 2020.
  * Boyd and Vandenberghe [2004] S. Boyd and L. Vandenberghe.  _Convex Optimization_.  Cambridge University Press, 2004.  doi: 10.1017/CBO9780511804441.
  * Campbell et al. [2020] D. Campbell, L. Liu, and S. Gould.  Solving the Blind Perspective-n-Point Problem End-To-End With Robust Differentiable Geometric Optimization.  In _Proc. of the European Conference on Computer Vision (ECCV)_ , 2020\.
  * Cao and Xie [2022] K. Cao and L. Xie.  Game-Theoretic Inverse Reinforcement Learning: A Differential Pontryagin’s Maximum Principle Approach.  _IEEE Transactions on Neural Networks and Learning Systems_ , pages 1–8, 2022.  doi: 10.1109/TNNLS.2022.3148376.
  * Cherian et al. [2017] A. Cherian, B. Fernando, M. Harandi, and S. Gould.  Generalized rank pooling for activity recognition.  In _Proc. of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2017.
  * Englert et al. [2017] P. Englert, N. A. Vien, and M. Toussaint.  Inverse KKT: Learning cost functions of manipulation tasks from demonstrations.  _The International Journal of Robotics Research_ , 36(13-14):1474–1488, 2017.  doi: 10.1177/0278364917745980.
  * Gill et al. [2005] P. E. Gill, W. Murray, and M. A. Saunders.  SNOPT: An SQP Algorithm for Large-Scale Constrained Optimization.  _SIAM Review_ , page 99–131, 2005.
  * Gould et al. [2022] S. Gould, R. Hartley, and D. Campbell.  Deep Declarative Networks.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 44:3988–4004, Aug 2022.  doi: 10.1109/TPAMI.2021.3059462.
  * Hajj and Skelboe [1990] I. N. Hajj and S. Skelboe.  A multilevel parallel solver for block tridiagonal and banded linear systems.  _Parallel Computing_ , 15:21–45, 1990.
  * Harris et al. [2020] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant.  Array programming with NumPy.  _Nature_ , 585:357–362, 2020.
  * Hatz et al. [2012] K. Hatz, J. P. Schlöder, and H. G. Bock.  Estimating parameters in optimal control problems.  _SIAM Journal on Scientific Computing_ , 34(3):A1707–A1728, 2012.  doi: 10.1137/110823390.
  * Hirshman et al. [2010] S. Hirshman, K. Perumalla, V. Lynch, and R. Sanchez.  Bcyclic: A parallel block tridiagonal matrix cyclic solver.  _Journal of Computational Physics_ , 229:6392–6404, 2010\.
  * Howell et al. [2019] T. A. Howell, B. E. Jackson, and Z. Manchester.  ALTRO: A Fast Solver for Constrained Trajectory Optimization.  In _Proc. of the IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pages 7674–7679, 2019.
  * Jin et al. [2020] W. Jin, Z. Wang, Z. Yang, and S. Mou.  Pontryagin Differentiable Programming: An End-to-End Learning and Control Framework.  _Advances in Neural Information Processing Systems_ , 33:7979–7992, 2020.
  * Jin et al. [2021a] W. Jin, D. Kulić, S. Mou, and S. Hirche.  Inverse optimal control from incomplete trajectory observations.  _The International Journal of Robotics Research_ , 40(6-7):848–865, 2021a.  doi: 10.1177/0278364921996384.
  * Jin et al. [2021b] W. Jin, S. Mou, and G. J. Pappas.  Safe Pontryagin Differentiable Programming.  _Advances in Neural Information Processing Systems_ , 34:16034–16050, 2021b.
  * Johnson et al. [2013] M. Johnson, N. Aghasadeghi, and T. Bretl.  Inverse optimal control for deterministic continuous-time nonlinear systems.  In _Decision and Control (CDC), 2013 IEEE 52nd Annual Conference on_ , pages 2906–2913, Dec 2013.
  * Keshavarz et al. [2011] A. Keshavarz, Y. Wang, and S. Boyd.  Imputing a convex objective function.  In _Intelligent Control (ISIC), 2011 IEEE International Symposium on_ , pages 613–619. IEEE, 2011.
  * Laine and Tomlin [2019] F. Laine and C. Tomlin.  Efficient computation of feedback control for equality-constrained lqr.  In _Proc. of the IEEE International Conference on Robotics and Automation (ICRA)_ , 2019.  doi: 10.1109/ICRA.2019.8793566.
  * Landry et al. [2019] B. Landry, Z. Manchester, and M. Pavone.  A differentiable augmented lagrangian method for bilevel nonlinear optimization.  In _Robotics: Science and Systems_ , 2019.
  * László et al. [2016] E. László, M. Giles, and J. Appleyard.  Manycore algorithms for batch scalar and block tridiagonal solvers.  _ACM Transactions on Mathematical Software_ , 42(4), 2016\.
  * Lee et al. [2019] K. Lee, S. Maji, A. Ravichandran, and S. Soatto.  Meta-learning with differentiable convex optimization.  In _Proc. of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2019.
  * Ljung [1999] L. Ljung.  _System identification: theory for the user_.  Prentice Hall PTR, Upper Saddle River, NJ, 1999.
  * Mastalli et al. [2020] C. Mastalli, R. Budhiraja, W. Merkt, G. Saurel, B. Hammoud, M. Naveau, J. Carpentier, L. Righetti, S. Vijayakumar, and N. Mansard.  Crocoddyl: An Efficient and Versatile Framework for Multi-Contact Optimal Control.  In _Proc. of the IEEE International Conference on Robotics and Automation (ICRA)_ , 2020.
  * Mensch and Blondel [2018] A. Mensch and M. Blondel.  Differentiable dynamic programming for structured prediction and attention.  In _Proc. of the International Conference on Machine Learning (ICML)_ , 2018.
  * Molloy et al. [2018] T. L. Molloy, J. J. Ford, and T. Perez.  Finite-horizon inverse optimal control for discrete-time nonlinear systems.  _Automatica_ , 87:442 – 446, 2018.  doi: 10.1016/j.automatica.2017.09.023.
  * Molloy et al. [2022] T. L. Molloy, J. Inga Charaja, S. Hohmann, and T. Perez.  _Inverse Optimal Control and Inverse Noncooperative Dynamic Game Theory: A Minimum-Principle Approach_.  Springer International Publishing, 2022.
  * Mombaur et al. [2010] K. Mombaur, A. Truong, and J.-P. Laumond.  From human to humanoid locomotion—an inverse optimal control approach.  _Autonomous robots_ , 28(3):369–383, 2010.
  * Pineda et al. [2022] L. Pineda, T. Fan, M. Monge, S. Venkataraman, P. Sodhi, R. T. Chen, J. Ortiz, D. DeTone, A. Wang, S. Anderson, J. Dong, B. Amos, and M. Mukadam.  Theseus: A Library for Differentiable Nonlinear Optimization.  _Advances in Neural Information Processing Systems_ , 2022.
  * Polizzi and Sameh [2006] E. Polizzi and A. H. Sameh.  A parallel hybrid banded system solver: the spike algorithm.  _Parallel Computing_ , 32(2):177–194, 2006.
  * Ratliff et al. [2006] N. D. Ratliff, J. A. Bagnell, and M. A. Zinkevich.  Maximum margin planning.  In _Proceedings of the 23rd international conference on Machine learning_ , pages 729–736. ACM, 2006.
  * Romero et al. [2023] A. Romero, Y. Song, and D. Scaramuzza.  Actor-critic model predictive control, 2023.
  * Russell et al. [2019] C. Russell, M. Toso, and N. Campbell.  Fixing implicit derivatives: Trust-region based learning of continuous energy functions.  In _Advances in Neural Information Processing Systems_ , volume 32, 2019.
  * Sarlin et al. [2020] P.-E. Sarlin, D. DeTone, T. Malisiewicz, and A. Rabinovich.  Superglue: Learning feature matching with graph neural networks.  In _Proc. of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2020.
  * Seal et al. [2013] S. K. Seal, K. S. Perumalla, and S. P. Hirshman.  Revisiting parallel cyclic reduction and parallel prefix-based algorithms for block tridiagonal system of equations.  _Journal of Parallel and Distributed Computing_ , 73, 2013.
  * Sodhi et al. [2021] P. Sodhi, E. Dexheimer, M. Mukadam, S. Anderson, and M. Kaess.  LEO: Learning energy-based models in factor graph optimization.  In _Proc. of the Conference on Robot Learning (CoRL)_ , 2021.
  * Tedrake [2023] R. Tedrake.  _Underactuated Robotics_.  2023\.  URL https://underactuated.csail.mit.edu.
  * Vlastelica et al. [2020] M. Vlastelica, A. Paulus, V. Musil, G. Martius, and M. Rolínek.  Differentiation of blackbox combinatorial solvers.  In _Proc. of the International Conference on Learning Representations (ICLR)_ , 2020.  *Equal Contribution.
  * Wächter and Biegler [2006] A. Wächter and L. Biegler.  On the implementation of an interior-point filter line-search algorithm for large-scale nonlinear programming.  _Mathematical Programming_ , 106:25–57, 2006.
  * Xiao et al. [2023] W. Xiao, T.-H. Wang, R. Hasani, M. Chahine, A. Amini, X. Li, and D. Rus.  Barriernet: Differentiable control barrier functions for learning of safe robot control.  _IEEE Transactions on Robotics_ , 39:2289–2307, 2023.
  * Xu et al. [2023] M. Xu, S. Garg, M. Milford, and S. Gould.  Deep declarative dynamic time warping for end-to-end learning of alignment paths.  In _Proc. of the International Conference on Learning Representations (ICLR)_ , 2023.
  * Yang et al. [2023] F. Yang, C. Wang, C. Cadena, and M. Hutter.  iPlanner: Imperative Path Planning.  In _Proceedings of Robotics: Science and Systems_ , Daegu, Republic of Korea, July 2023.  doi: 10.15607/RSS.2023.XIX.064.
  * Yi et al. [2021] B. Yi, M. Lee, A. Kloss, R. Martín-Martín, and J. Bohg.  Differentiable factor graph optimization for learning smoothers.  In _Proc. of the IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , 2021.
  * Zanon and Gros [2021] M. Zanon and S. Gros.  Safe Reinforcement Learning Using Robust MPC.  _IEEE Transactions on Automatic Control_ , 66:3638–3652, 2021.
  * Ziebart et al. [2008] B. D. Ziebart, A. L. Maas, J. A. Bagnell, and A. K. Dey.  Maximum entropy inverse reinforcement learning.  In _AAAI Conference on Artificial Intelligence_ , pages 1433–1438, 2008.

##  Appendix A Experimental Setup for Imitation Learning/LfD Task

We use the code and simulation environments provided by [22] and follow a
similar experimental setup, described in more detail in this section. The LfD
experimental setting without inequality constraints is referred to in [20] as
the inverse reinforcement learning (IRL) task, whereas including inequality
constraints is referred to in [22] as the constrained inverse optimal control
(CIOC) task.

#### Demonstration Trajectories.

Up to five demonstration trajectories are used in the LfD setting with no
inequality constraints. Each demonstration trajectory is generated by solving
a COC problem using the same underlying parameters $\theta$, however the
initial conditions differ across trajectories and are sampled randomly. For
the setting with inequality constraints present, we use only one demonstration
trajectory.

#### Initialization.

Consistent with [20, 22], we run five imitation learning trials for all LfD
experiments, where for each trial, $\theta$ is initialized by adding uniform
noise to the true value. For all methods, we use the gradient descent with the
same learning rate for a given environment.

#### Environment Specifications.

Table 1 provides summary specifications of the simulation environments used in
our experiments, which are consistent with prior works [20, 22]. Recall that
$n$, $m$ denote the number of states and controls, respectively. The number of
parameters is denoted $d=|\theta|$ and $T$ is the horizon length.

Table 1: Description of environments Environment | $n$ | $m$ | $d$ | $T_{\text{IRL}}$ | $T_{\text{CIOC}}$
---|---|---|---|---|---
Cartpole | 4 | 1 | 9 | 30 | 35
Quadrotor | 13 | 4 | 11 | 50 | 25
Robotarm | 4 | 2 | 10 | 35 | 25
Rocket | 13 | 3 | 12 | 40 | 40

#### Additional Environment Parameters.

Additional specifications are presented in Table 2. These include log-barrier
parameter $\gamma$ used for Safe-PDP (b), time discretization for Forward
Euler to discretize continuous time dynamics, the learning rate for minimizing
the imitation loss and finally the number of demonstration trajectories per
environment. (E) refers to experiments without constraints, whereas (I) means
inequality constraints are present.

Table 2: Additional hyperparameters for LfD experiments Environment | $\gamma$ | $\Delta$ | lr | $n_{\text{demos}}$
---|---|---|---|---
Cartpole (E) | - | 0.1 | $10^{-4}$ | 5
Quadrotor (E) | - | 0.1 | $10^{-4}$ | 2
Robotarm (E) | - | 0.1 | $10^{-4}$ | 4
Rocket (E) | - | 0.1 | $3\times 10^{-4}$ | 1
Cartpole (I) | 0.01 | 0.1 | $8\times 10^{-5}$ | 1
Quadrotor (I) | 0.01 | 0.15 | $2\times 10^{-4}$ | 1
Robotarm (I) | 0.01 | 0.2 | $2\times 10^{-3}$ | 1
Rocket (I) | 1 | 0.1 | $10^{-5}$ | 1

###  A.1 Cartpole

The cartpole environment relates to the swing-up and balance task, where a
simple pendulum (i.e., massless pole and point mass on the end of the pole) is
swinging on a cart and the objective is to balance the pendulum above the
cart. Only a horizontal force can be applied to drive the cart.

#### Dynamics.

The state variable $x=[y,q,\dot{y},\dot{q}]^{\top}\in\mathbb{R}^{4}$ is
comprised of the horizontal position of the cart $y$, counter-clockwise angle
of the pendulum $\theta$ (from the hanging position) and their respective
velocities. Control $u\in\mathbb{R}$ relates to the horizontal force applied
to the cart. Dynamics parameters
$\theta_{\text{dyn}}=[m_{c},m_{l}p,\ell]^{\top}\in\mathbb{R}^{3}$ relate to
the the mass of the cart $m_{c}$, mass of the point mass at the end of the
pole $m_{p}$ and length of the pole $\ell$. Gravity is set at $g=10$. The
continuous time dynamics are given by

$\displaystyle\ddot{y}$
$\displaystyle=\frac{1}{m_{c}+m_{c}\sin^{2}q}\left[u+m_{p}\sin
q(\ell\dot{q}^{2}+g\cos q)\right]$ (8) $\displaystyle\ddot{q}$
$\displaystyle=\frac{1}{\ell(m_{c}+m_{p}\sin^{2}q)}\left[-u\cos
q-m_{p}\ell\dot{q}^{2}\cos q\sin q-(m_{c}+m_{p})g\sin q\right],$ (9)

with a detailed derivation provided in [43, Ch. 3].

#### Cost Functions.

The cost function is a quadratic cost over states and controls. Specifically,

$c_{t}(x_{t},u_{t};\theta)=(x_{t}-x_{\text{goal}})^{\top}W(x_{t}-x_{\text{goal}})+w_{u}u^{2},$
(10)

and

$c_{T}(x_{T},u_{T};\theta)=(x_{T}-x_{\text{goal}})^{\top}W(x_{T}-x_{\text{goal}}),$
(11)

with $x_{\text{goal}}=[0,0,-\pi,0]^{\top}$. We specify that
$W\triangleq\text{diag}(w)$, where
$w=[w_{y},w_{q},w_{\dot{y}},w_{\dot{q}}]\geq 0$ and furthermore, $w_{u}\geq
0$. We have that
$\theta_{\text{cost}}=[w^{\top},w_{u}^{\top}]^{\top}\in\mathbb{R}^{5}$.

#### Constraints.

For the setting with constraints, we apply box constraints for the cart
position and controls. Specifically, we enforce $|y|\leq y_{\text{max}}$ and
$|u|\leq u_{\text{max}}$. We let constraint parameters
$\theta_{\text{constr}}=[y_{\text{max}},u_{\text{max}}]^{\top}$. The final
parameter vector is given by
$\theta=[\theta_{\text{dyn}}^{\top},\theta_{\text{cost}}^{\top},\theta_{\text{constr}}^{\top}]^{\top}$.

###  A.2 6-DoF Quadrotor Manouvering

See Section E of the appendix in [20] for the full definition of the quadrotor
maneuvering problem. The objective is to drive the quadrotor using propeller
thrusts to a goal configuration. The state is given by
$[p^{\top},\dot{p}^{\top},q^{\top},\omega^{\top}]^{\top}\in\mathbb{R}^{13}$,
where $p\in\mathbb{R}^{3}$ represents the position of the center-of-mass of
the quadrotor, $\dot{p}\in\mathbb{R}^{3}$ represents linear velocity,
$q\in\mathbb{R}^{4}$ is a unit quaternion representation of the quadrotor
attitude and $\omega\in\mathbb{R}^{3}$ represents angular velocity. For the
inequality constrained setting, non-linear (squared-norm) constraints on
quadrotor position, given by $\|p\|_{2}^{2}\leq r_{\text{max}}$ and thrust
limits are applied.

###  A.3 2-Link Robot Arm

See Section E in the appendix in [20] for a reference to the full description
of the robot arm problem. The objective of this OC problem is to move the
robot arm to a goal configuration. The state is given by the orientation of
the base and elbow joint, as well as respective velocities
$[q_{1},q_{2},\dot{q}_{1},\dot{q}_{2}]^{\top}\in\mathbb{R}^{4}$. The control
inputs $u=[u_{q},u_{2}]^{\top}\in\mathbb{R}^{2}$ relate to applying torques at
each joint. Constraints for the inequality constrained setting relate to box
constraints on the joint positions, as well as torque limits. The cost
function is similar to cartpole, i.e., a quadratic penalty to a goal state and
over the controls.

###  A.4 6-DoF Rocket Landing

See Section I in the appendix in [20] for the full definition of the 6-DoF
rocket landing problem. The objective of this task is to land a rocket modeled
as a rigid body (softly) onto a goal position, using thrusters located at the
tail. A non-linear constraint over the tilt angle is applied in the inequality
constrained setting. The cost function penalizes deviations from the goal
configuration, as well as fuel cost.

##  Appendix B Experimental Setup for Synthetic Benchmark

We discuss in more detail how blocks for $H,A,B,C$ are generated in this
section. Elements for the Jacobian and Hessian blocks for the objective
function and constraints are sampled independently and identically distributed
(i.i.d.) from a standard Gaussian. Hessian blocks are made symmetric using
$H^{\text{sym}}=(H+H^{\top})/2$, where $H$ is the initially generated random
matrix. Each block’s condition number is modified by applying an SVD to each
block and adjusting the diagonal entries. Elements of the downstream loss
gradients $v=\text{D}_{\xi}\mathcal{L}(\xi)^{\top}$ are also sampled i.i.d.
from a standard Gaussian.

##  Appendix C Additional Results for the Imitation Learning/LfD Tasks

In this section, we provide additional detailed analysis around the imitation
learning/LfD setting with no inequality constraints. Figure 6 illustrates the
results of these experiments. We plot the percentage difference in loss across
the full training curve between IDOC and PDP for five random initializations
for $\theta$. For the quadrotor, robotarm and rocket environments, we observe
almost no difference in learning curves. This is expected since IDOC and PDP
are computing the same trajectory derivative. For cartpole however, we see
that PDP fails catastrophically (resulting in undesirable spikes in the
imitation loss) on two occasions across the five trials. We suspect this is
due to numerical instability of the Riccati equations, arising from the local
geometric properties of the dynamics at solved trajectories.

(a) Cartpole

(b) 6-DoF Quadrotor

(c) Robotarm

(d) 6-DoF Rocket

Figure 6: Results for the imitation learning task with no inequality
constraints present. The y-axis is the (percentage) relative difference in the
imitation loss, defined to be (IDOC loss - PDP loss) / PDP loss (lower is
better). For Cartpole, locations where the PDP gradient fails is marked by
stars. Shaded regions indicate min/max values over five trials.

##  Appendix D Additional Timings for Simulation Experiments

In this section, we provide additional compute time experiments for the
imitation learning/LfD tasks evaluated in the simulation environments. Given
the relatively small size for the COC problems, there is very limited
advantage with respect to compute time for parallel evaluation of trajectory
derivatives. However, as a proxy to parallelization, IDOC can be vectorized
across timesteps using the numerical linear algebra library Numpy [16]. All
experiments in this section are run on a single thread of an AMD Ryzen 9 7900X
4.7Ghz desktop CPU.

The experimental setting where inequality constraints are present is
challenging to vectorize compared to the setting with only equality
constraints, because different timesteps may have different numbers of active
inequality constraints. While matrices $A$ and $C$ presented in Equation 4 are
still block structured, the blocks may not be of uniform size in this setting.
We present in Figure 7(a) and 7(b) the mean computation time (along with the
upper/lower bound) per iteration across five trials with 10k iterations for
each trial for all environments except for rocket (1k). We observe around a
2$\times$ speedup over PDP by computing our IFT/DDN gradients, even before
direct computation of vector-Jacobian products. Computing VJPs directly yields
a further (approx.) 10% saving in compute time over constructing the
trajectory derivative explicitly.

In our non-optimized Python implementation of IDOC, we batch together and
vectorize computations involving blocks with an identical number of active
constraints. As expected, trajectory derivatives for inequality constrained
problems are slightly slower to compute compared to the equivalent without
(hard) inequality constraints (log-barrier approximation), due to the
necessary computational overheads required for identifying and batching
blocks.

Note that Safe-PDP [22] (for hard inequality constraints) and PDP [20] (used
for differentiating through log-barrier problems) derivatives are implemented
using a custom LQR and equality constrained LQR solver, respectively. These
solvers have markedly different implementations, which explains differences in
computation time between Safe-PDP and PDP.

(a) Hard Constraints

(b) Log-Barrier

Figure 7: Computation times for trajectory derivatives (per trajectory, per
iteration) for the CIOC task. As expected, the log-barrier method with no
“hard" inequality constraints affords faster derivative computation.