biblatexFile ’english-ieee.lbx’ soft open fences parencite()

B.E.(Hons.), BITS Pilani Master of Computing 2024 Assistant Professor Harold
Soh

Assistant Professor Bryan Hooi

# Planning using Schrödinger Bridge Diffusion Models

Adarsh Srivastava

###### Abstract

Offline planning often struggles with poor sampling efficiency as it tries to
learn policies from scratch. Especially with diffusion models, such cold start
practices mean that both training and sampling become very expensive. We
hypothesize that certain environment constraint priors or cheaply available
policies make it unnecessary to learn from scratch, and explore a way to
incorporate such priors in the learning process. To achieve that, we borrow a
variation of Schrödinger bridge formulation from image-to-image setting and
apply it to planning tasks. We study the performance on some planning tasks
and compare the performance against the DDPM formulation.

23 May 2024 pic/signature.png

###### Acknowledgements.

I would like to thank Prof. Harold Soh for giving me this opportunity, and for
his time, guidance and patience with me throughout the thesis. I am also
grateful to Alvin Heng, a Ph.D. student under Prof. Soh, with whom I had many
fruitful discussions and who patiently walked me through several difficult
concepts, and whose discussions with Prof. Soh helped spark many ideas,
including the one for this thesis. Simply listening to his discussions has
allowed me to learn much and think in newer directions. I am also grateful to
my friends in NUS and beyond, whose constant encouragement kept me going.

###### Contents

  1. 0 Introduction
  2. 1 Background & Literature Review
    1. 1 Diffusion models
      1. 1 Diffusion Probabilistic models
      2. 2 Score-Based Generative Modeling through SDEs
    2. 2 Diffusion models in Planning
    3. 3 The Diffuser approach
      1. 1 Architecture
      2. 2 Key Results
    4. 4 Schrödinger Bridges
      1. 1 Schrödinger Bridge Problem
      2. 2 I2SB: Image-to-Image Schrödinger Bridge
        1. Original Schrödinger Bridge SDEs
        2. From SB to I2SB
        3. Implementation
        4. Key Results
  3. 2 Methodology
  4. 3 Evaluation
    1. 1 Experimental Setup
      1. 1 Environment and Task
      2. 2 Task Details and Metrics
      3. 3 Prior Construction
        1. Analytical Prior
        2. Learned Prior
        3. Random Prior
    2. 2 Results
      1. 1 Uninformative Random Prior
        1. Diffusing from Gaussian noise - Same NFE
        2. Diffusing from Gaussian noise - across a range of NFEs
      2. 2 Informative Priors
    3. 3 Summary of Experimental Results
      1. 1 Sample Efficiency
      2. 2 Sampling Efficiency
      3. 3 On Priors
      4. 4 On the Performance Gap between I2SB and DDPM
  5. 4 Discussion & Conclusion

##  Chapter 0 Introduction

Planning from scratch has long struggled with poor sample efficiency and
therefore the need for a large amount of data and compute. Learning from
scratch is not only expensive, but also sometimes unnecessary, either because
natural domain or task constraints already exist, or because an approximate
policy is cheaply available. Not being able to incorporate these into the
planning algorithm naturally not only leads to higher data requirements ( or
poor sample efficiency) but also longer training and sampling times, since the
model needs to learn information that is readily available and can be supplied
as a prior.

Prior methods that have tried to deal with this issue have either tried to
incorporate learned skills as priors between tasks [13], using heuristics to
accelerate the algorithm [4], or warm starting policies [21].

In this thesis, we explore a method of planning that directly incorporates a
prior policy into the training, such that learning starts from the prior
policy distribution itself. To do so, we build upon two prior works - using
diffusion models to learn trajectory planning and optimization as first
described in [9], and a bridge diffusion model that transports one
distribution to another, formulated for image-to-image tasks such as
inpainting and deblurring in [11].

Diffusion models [17] [8] [18] have shown state-of-the-art results in many
image and video generation tasks including unconditioned generation, text-
conditioned generation, image inpainting, superresolution, and Image-to-Image
Translation tasks [12]. In [9], they have also shown their ability to capture
complex target distributions and ushered in a new paradigm of learned planning
- one where planning and trajectory optimization were wrapped into a single
learned model. However, while diffusion models beat other generative models in
high-quality sample generation capabilities with mode coverage and diversity
as well as stable training, they also typically suffer from expensive sampling
and long training times.

Figure 1: Diffusion models suffer from expensive sampling and long training
times.

A key reason for this inefficiency is that diffusion models learn to denoise
to the target distribution from pure Gaussian noise. Intuitively, if it was
possible to sample from a distribution that was closer to the target
distribution than pure noise (in other words, it had some structural
information about the target), it should make sense that a similar process as
diffusion would have less to learn ( and therefore learn better from fewer
samples, i.e., improved sample efficiency), and would be less expensive to
sample from (since the generation has less to predict).

This idea was explored by [11] in the context of image-to-image translation
tasks using diffusion bridges built upon the theoretical framework of the
Schrödinger Bridge problem, but so far has not been applied to planning or
imitation learning tasks. We hypothesize that using a bridging method to
generate plans should also work, and should lead to better sample efficiency
and faster sampling.

Since I2SB [11] and Diffuser [9] share the common framework of diffusion
models, and I2SB shows good performance on image translation tasks, we pick
I2SB as the choice of bridging model to explore how effective a pairing of the
two approaches can be in a planning setting.

This thesis is structured as follows: we first go through the necessary
background on diffusion models, the Diffuser approach, the Schrödinger Bridge
problem, the I2SB version of Schrödinger Bridges, and their adaptation to
image-to-image tasks. We then outline our own approach and lay out our
experimental setup. Finally, we discuss the results and future directions.

We show that while the tractable class of Schrödinger Bridge we use to bridge
distributions works, and shows better performance than DDPM at very low NFEs,
at higher NFEs DDPM catches up and is able to learn and generalize better. We
also show how we might construct priors for tasks in planning for such
bridging methods. However, this work does not go into the theoretical analysis
of the various bridging models, as that is beyond the scope of this work.

##  Chapter 1 Background & Literature Review

This chapter outlines the theoretical foundations which this thesis builds
upon. First, we give a brief overview of diffusion models in two forms - the
probabilistic model [8], and the score-based model [19]. We also briefly cover
the use of diffusion models in planning, and in particular, the Diffuser
approach [9], which we use as the foundation and baseline for our approach. We
then cover the theoretical foundations at the crux of our approach - the
Schrödinger Bridge, and one particular formulation of it - I2SB [11], which is
the formulation we use in our experiments.

###  1 Diffusion models

Diffusion models are a class of deep generative models inspired by non-
equilibrium thermodynamics that iteratively corrupt data with increasing
noise, and learn to reverse this process from a noise sample to form a
generative model for the data distribution. Denoising diffusion probabilistic
modeling (DDPM) [17] [8] learns a denoising Markov chain process with discrete
time steps, while score-based models [19] learn the gradient of the log data
distribution, and use it to guide a random towards regions of higher data
density. In particular, score-based modeling through SDEs [20] models the
diffusion process through continuous-time forward and reverse SDEs.

####  1 Diffusion Probabilistic models

A T-step Denoising Diffusion Probabilistic Model(DDPM) consists of two
processes: the forward process (also referred to as the diffusion process),
and the reverse inference process.

Figure 1: The DDPM graphical model.

The forward process, or the diffusion process, adds Gaussian noise to the
sample according to a fixed variance schedule $\beta_{1},\dotsc,\beta_{T}$:

$q\left(\mathbf{x}_{1:T}\mid\mathbf{x}_{0}\right):=\prod_{t=1}^{T}q\left(\mathbf{x}_{t}\mid\mathbf{x}_{t-1}\right),\quad
q\left(\mathbf{x}_{t}\mid\mathbf{x}_{t-1}\right):=\mathcal{N}\left(\mathbf{x}_{t};\sqrt{1-\beta_{t}}\mathbf{x}_{t-1},\beta_{t}\mathbf{I}\right)$
(1)

There is also a closed-form expression for sampling the forward process at any
$x_{t}$; using $\alpha_{t}:=1-\beta_{t}$ and
$\bar{\alpha}_{t}:=\prod_{s=1}^{t}\alpha_{s}$, we have:

$q\left(\mathbf{x}_{t}\mid\mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t};\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right)\mathbf{I}\right)$
(2)

The reverse process, or the denoising process, is defined as a Markov chain
with learned Gaussian transitions starting at
$p\left(\mathbf{x}_{T}\right)=\mathcal{N}\left(\mathbf{x}_{T};\mathbf{0},\mathbf{I}\right)$:

$p_{\theta}\left(\mathbf{x}_{0:T}\right):=p\left(\mathbf{x}_{T}\right)\prod_{t=1}^{T}p_{\theta}\left(\mathbf{x}_{t-1}\mid\mathbf{x}_{t}\right),\quad
p_{\theta}\left(\mathbf{x}_{t-1}\mid\mathbf{x}_{t}\right):=\mathcal{N}\left(\mathbf{x}_{t-1};\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t},t\right),\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t},t\right)\right)$
(3)

The model learns by training a time-dependent UNet-like network to predict the
total noise, $\boldsymbol{\epsilon}_{\theta}$, added to the sample during the
forward process at any timestep $t$:

$L_{\text{simple
}}(\theta):=\mathbb{E}_{t,\mathbf{x}_{0},\boldsymbol{\epsilon}}\left[\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\theta}\left(\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}}\boldsymbol{\epsilon},t\right)\right\|^{2}\right]$
(4)

####  2 Score-Based Generative Modeling through SDEs

Another approach [20] treats the forward process as a continuous time-
dependent stochastic differential equation (SDE) given by ($\mathbf{w}$ is the
standard Wiener process or Brownian motion):

$\mathrm{d}\mathbf{x}=\mathbf{f}(\mathbf{x},t)\mathrm{d}t+g(t)\mathrm{d}\mathbf{w}$
(5)

and the reverse process reverse-time SDE by:

$\mathrm{d}\mathbf{x}=\left[\mathbf{f}(\mathbf{x},t)-g(t)^{2}\nabla_{\mathbf{x}}\log
p_{t}(\mathbf{x})\right]\mathrm{d}t+g(t)\mathrm{d}\overline{\mathbf{w}}$ (6)

Figure 2: Overview of score-based generative modeling through SDEs.

The model learns by training a time-dependent UNet-like network to predict the
score $\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})$ of the data distribution for
the marginal distribution $p(x)$ and continuous time $t$. With this score
available, sampling is done by integrating the reverse SDE using numerical SDE
solvers.

[20] also talk about two kinds of SDEs - the Variance Exploding (VE) SDE, and
the Variance Preserving (VP) SDE, depending on whether the variance of the
forward SDe explodes as $t\rightarrow\infty$ or remain bounded. In particular,
they note that SMLD [19] is a VE model, and DDPM [8] is a VP model.

###  2 Diffusion models in Planning

Diffusion models trained to generate images or videos have previously been
used in the planning domain in several works. DALL-E-Bot [10] uses the DALL-E
[14] foundation model trained on web-scale data to generate text-conditioned
goal images for object rearrangement in a scene. UniPi [6] uses text-to-video
diffusion models to construct a video plan of how to achieve a task, followed
by an inverse dynamics model to extract control actions. ROSIE [22] uses
conditional inpainting of the current scene using text-to-image diffusion
models to create variations and improve generalization during training without
needing more real-world data.

However, one of the earliest works that directly trained diffusion models to
capture the underlying trajectory/policy distribution was Diffuser [9] \-
which used diffusion over state-action pairs to generate viable variable
horizon trajectories. This work showed that diffusion models can generate
optimized trajectories non-autoregressively over long horizons even for
complex tasks, without suffering from compounding errors, and can also compose
in-distribution trajectories to generate novel trajectories. Since this thesis
builds upon the Diffuser approach, we will explore the technical details of
this approach in depth in the following section.

###  3 The Diffuser approach

Traditionally, planning was accomplished by learning a dynamics model of the
environment and then using trajectory optimizers to generate the trajectory.
The key idea behind the Diffuser [9] approach was to wrap the two into a
single diffusion model that learns to generate optimized trajectories for the
environment and given constraints.

####  1 Architecture

To reap the advantages of work done on image-based diffusion models, Diffuser
treats trajectories it trains on much like an image - the state and action
values are analogous to the pixel values, the timesteps being the length, and
the width being the sum of dimensionalities of the state and action values.
Specifically, the input (and output) trajectories $\tau$ are defined as the
following array, $T$ being the prediction horizon.

$\boldsymbol{\tau}=\left[\begin{array}[]{llll}\mathbf{s}_{0}&\mathbf{s}_{1}&\ldots&\mathbf{s}_{T}\\\
\mathbf{a}_{0}&\mathbf{a}_{1}&\ldots&\mathbf{a}_{T}\end{array}\right]$

Because planning and trajectory optimization are wrapped into one, and
trajectory optimization is inherently non-markovian (the present state depends
on both the past and future states), Diffuser predicts (or denoises) all
timesteps at once, instead of predicting autoregressively.

Despite the non-autoregressive prediction, Diffuser maintains temporal local
consistency by using a small local receptive field along the planning horizon,
implemented using 1D convolutions over the time dimension. Also, because
Diffuser uses these temporal convolutions, the architecture is independent of
the planning horizon to allow variable-length plans.

Figure 3: Diffuser uses a local receptive field to enforce local consistency
while predicting non-autoregressively.

Diffuser also allows for a guidance function $\mathcal{J}$ to influence the
generated trajectory. This guidance function can represent prior knowledge,
rewards, or costs associated with the trajectory.

Because Diffuser uses a temporal local receptive field to enforce local
consistency, an interesting way to generate a goal-seeking policy is to treat
the problem as an inpainting problem, where the start and end states are kept
fixed, and the model denoises the remaining trajectory. This is the method
used in the Maze2D benchmark experiments run by the authors.

Figure 4: Diffuser treats generating a goal-seeking policy as an inpainting
problem.

The overall architecture of Diffuser remains remarkably similar to image
diffusion models, except that it replaces the spatial convolution UNet with a
temporal convolution UNet, with 1D convolutions over the planning horizon.

####  2 Key Results

The authors show the following key aspects of planning with Diffuser:

  * •

Ability to perform long horizon planning, because Diffuser does not suffer
from compounding errors, and losses are calculated over the whole trajectory
instead of minimizing single step error.

  * •

Ability to generate variable length plans, since the architecture is
independent of the planning horizon.

  * •

Ability to compose in-distribution trajectories to generate novel
trajectories.

Figure 5: A visualization of non-autoregressive denoising of trajectories on
Maze2D tasks.

Diffuser showed competitive results on D4RL benchmarks across tasks compared
to other state-of-the-art algorithms, once again confirming the ability of
Diffusion models to capture complex distributions with relative ease. A
limitation of Diffuser is that individual plans are slow to generate due to
the iterative nature of the algorithm.

###  4 Schrödinger Bridges

Directly transporting between two distributions is an important problem,
especially in image-to-image translation tasks. A flurry of recent work has
happened on coming up with ways to achieve this transport. [1] introduces
Stochastic Interpolants that bridge two arbitrary densities exactly and in
finite time, which can be used for image generation and translation. [23]
propose a unified diffusion bridge method as a natural extension of DDPM [8].

For this thesis, we borrow the formulation of a tractable class of Schrödinger
Bridge from the I2SB paper[11], which used the technique on image-to-image
translation tasks such as inpainting, super-resolution, and deblurring. In the
following sub-sections, we go through the details of their methods, and also
lay down some theoretical background.

####  1 Schrödinger Bridge Problem

The Schrödinger Bridge Problem [16] is finding the most likely evolution of a
stochastic process between two continuous distributions, and is a classical
problem appearing in areas of optimal control and probability. It can also be
seen as an entropy-regularized optimal transport problem.

Schrödinger Bridges (SB) can also be seen as the general case of a score-based
generative model (SGM), since while SGM degrades to and denoises from pure
noise, SB goes beyond pure noise to another distribution.

Despite the mathematical generalization, SGM and SB have evolved independently
and thus have different computational frameworks to solve them. Schrödinger
Bridge models often use iterative projection methods [3] which have
intractable complexity for high dimensional problems, and the methods are
different from how diffusion models are trained.

Figure 6: SB is 6x slower and 3x more memory intensive on high dimensional
inputs on similar tasks as SGM. [11]

####  2 I2SB: Image-to-Image Schrödinger Bridge

Diffusion bridges have been recently explored in [2], [3], [11], and [23], and
have shown to perform well on tasks like image-to-image translation, image
inpainting, image restoration, and super-resolution.

Since we use the same formulation of Schrödinger Bridge as the I2SB paper
[11], we describe below their approach and their theoretical and experimental
results.

The key idea behind the I2SB paper was that degraded images still contain a
lot of useful information about the final image, and thus diffusing to and
from Gaussian white noise, which has no structural information of the target
distribution, is bound to take longer to learn, and more diffusion steps to
infer. To better leverage the structure of the problem, I2SB directly starts
the generative process from the degraded image and builds diffusion bridges
between the two distributions. To construct such diffusion bridges, the
authors came up with a tractable class of Schrödinger Bridge that works for
this problem.

Figure 7: I2SB found a tractable way to directly bridge from degraded to clean
images without going to and from Gaussian noise.

In the following subsections, we discuss the theory behind the I2SB model, and
some of their results on image-to-image tasks.

##### Original Schrödinger Bridge SDEs

Given two boundary distributions $p_{\mathcal{A}}$ and $p_{\mathcal{B}}$ in
two distinct domains, with $X_{0}\sim p_{\mathcal{A}}$ and $X_{1}\sim
p_{\mathcal{B}}$, where $X_{t}$ is indexed by $t\in[0,1]$, the original
Schrödinger Bridge [16] model consider the following forward and backward
SDEs:

$\mathrm{d}X_{t}=\left[f_{t}+\beta_{t}\nabla\log\Psi\left(X_{t},t\right)\right]\mathrm{d}t+\sqrt{\beta_{t}}\mathrm{\leavevmode\nobreak\
d}W_{t}$ (7)

$\mathrm{d}X_{t}=\left[f_{t}-\beta_{t}\nabla\log\widehat{\Psi}\left(X_{t},t\right)\right]\mathrm{d}t+\sqrt{\beta_{t}}\mathrm{\leavevmode\nobreak\
d}\bar{W}_{t}$ (8)

The functions $\Psi,\widehat{\Psi}\in
C^{2,1}\left(\mathbb{R}^{d},[0,1]\right)$ are time-varying energy potentials
that solve the following coupled PDEs,

$\left\\{\begin{array}[]{l}\frac{\partial\Psi(x,t)}{\partial
t}=-\nabla\Psi^{\top}f-\frac{1}{2}\beta\Delta\Psi\\\
\frac{\partial\widehat{\Psi}(x,t)}{\partial
t}=-\nabla\cdot(\widehat{\Psi}f)+\frac{1}{2}\beta\Delta\widehat{\Psi}\end{array}\right.$
(9)

$\text{ s.t.
}\Psi(x,0)\widehat{\Psi}(x,0)=p_{\mathcal{A}}(x),\Psi(x,1)\widehat{\Psi}(x,1)=p_{\mathcal{B}}(x)$
(10)

It is this coupling in the above equations that gives general SB methods
unfavorable complexity on high dimensional tasks compared to SGM. (Figure 6)

From Equation 5 and Equation 7, we see the equations differ only in the
additional nonlinear drift terms $\nabla\log\Psi$. The forward drift
$\nabla\log\Psi$ is what allows the process to transport samples beyond
Gaussian priors. This also shows how SGMs are a special case of general SB.

#####  From SB to I2SB

To go from the general Schrödinger Bridge to a tractable version, the authors
make some modifications to the formulation. First, we observe that the
nonlinear drifts in Equation 7 resemble the score function in Equation 6 when
we view $\Psi(\cdot,t)$ and $\widehat{\Psi}(\cdot,t)$ as densities. Then, we
note that $\nabla\log\widehat{\Psi}\left(X_{t},t\right)$ and
$\nabla\log\Psi\left(X_{t},t\right)$ are the score functions of the following
linear SDEs respectively:

$\mathrm{d}X_{t}=f_{t}\left(X_{t}\right)\mathrm{d}t+\sqrt{\beta_{t}}\mathrm{\leavevmode\nobreak\
d}W_{t},\quad X_{0}\sim\widehat{\Psi}(\cdot,0)$ (11)

$\mathrm{d}X_{t}=f_{t}\left(X_{t}\right)\mathrm{d}t+\sqrt{\beta_{t}}\mathrm{\leavevmode\nobreak\
d}\bar{W}_{t},\quad X_{1}\sim\Psi(\cdot,1)$ (12)

Remember that we use neural nets to parameterize the score functions of linear
SDE in SGM (Equation 5). Therefore, if we can parameterize
$\nabla\log\widehat{\Psi}$ with the score network, we can use SGM to learn
$\nabla\log\widehat{\Psi}$. A similar logic goes for the second equation.

However, we still cannot sample from $\widehat{\Psi}(\cdot,0)$ or
$\Psi(\cdot,1)$ because of the coupling in Equation 10. So the authors present
a tractable case with a boundary modification.

If we let $p_{\mathcal{A}}(\cdot):=\delta_{a}(\cdot)$ be the Dirac Delta
distribution centered at some data point a, we have:

$\widehat{\Psi}(\cdot,0)=\delta_{a}(\cdot),\quad\Psi(\cdot,1)=\frac{p_{\mathcal{B}}}{\widehat{\Psi}(\cdot,1)}$
(13)

This is equivalent to the backward drift driving the reverse process of
Equation 11 always flows towards $a$, irrespective of $p_{\mathcal{B}}$.
Although this might mean that, in theory, we would not generalize beyond such
data points in the training sample, the authors rely on the strong
generalization abilities of neural networks to overcome this in practice.

##### Implementation

Based on the above modifications, the authors come up with a training and
sampling plan. Let $X_{0}$ be the target data point from
$p_{\mathcal{A}}\left(X_{0}\right)$ and $X_{1}$ be its degraded pair from
$p_{\mathcal{B}}\left(X_{1}\mid X_{0}\right)$.

Training: The authors derive the forward process posterior of Equation 7 to be
the following (derivation is deferred to the original paper):

$\begin{gathered}q\left(X_{t}\mid
X_{0},X_{1}\right)=\mathcal{N}\left(X_{t};\mu_{t}\left(X_{0},X_{1}\right),\Sigma_{t}\right)\\\
\mu_{t}=\frac{\bar{\sigma}_{t}^{2}}{\bar{\sigma}_{t}^{2}+\sigma_{t}^{2}}X_{0}+\frac{\sigma_{t}^{2}}{\bar{\sigma}_{t}^{2}+\sigma_{t}^{2}}X_{1},\quad\Sigma_{t}=\frac{\sigma_{t}^{2}\bar{\sigma}_{t}^{2}}{\bar{\sigma}_{t}^{2}+\sigma_{t}^{2}}\cdot
I\end{gathered}$ (14)

where $\sigma_{t}^{2}:=\int_{0}^{t}\beta_{\tau}\mathrm{d}\tau$ and
$\bar{\sigma}_{t}^{2}:=\int_{t}^{1}\beta_{\tau}\mathrm{d}\tau$ are variances
accumulated from either sides.

This is analogous in meaning and form to the closed form posterior during
training of DDPM in Equation 2.

Sampling: The sampling procedure for I2SB can be done exactly the same as that
of DDPM. However, the authors also show that the same formula used above
during training can also be used during sampling.

During sampling, typically, DDPM needs to recursively run each timestep in
Equation 3 one by one with the model generating the required
$\boldsymbol{\epsilon}_{\theta}$. However, in the I2SB formulation, the above
posterior also marginalizes the recursive sampling (the authors prove this
using induction utilizing Equation 14, proof is deferred to the original
paper) such that:

$q\left(X_{n}\mid X_{0},X_{N}\right)=\int\Pi_{k=n}^{N-1}p\left(X_{k}\mid
X_{0},X_{k+1}\right)\mathrm{d}X_{k+1}$ (15)

Therefore, we can sample multiple diffusion steps in one function call (1 NFE)
of the network, though this is not necessary.

##### Key Results

I2SB uses Frechet Inception Distance (FID) scores and classifier accuracy of a
pre-trained ResNet50 to compare against other image-to-image translation
algorithms, on 4 primary tasks - 4 x super-resolution, Inpainting, JPEG
restoration, and deblurring.

The key takeaway from those results is that I2SB surpasses or matches the
performance of the other algorithms, with a much greater sampling efficiency,
i.e., with far fewer function calls, which makes sense, since I2SB starts
generation from a much more structurally informative prior.

Figure 8: I2SB performs better than Palette [15] at lower NFEs, suggesting a
higher sampling efficiency. Figure 9: I2SB has much better FID and CA scores
at lower NFEs compared to Palette [15].

##  Chapter 2 Methodology

To test our hypothesis, we keep the same UNet network, architecture, and
structure as the Diffuser approach for planning with DDPM. We then replace the
DDPM training and sampling steps with the I2SB Schrödinger bridge steps using
Equation 14 and Equation 15 respectively.

There is an additional component to our setup - the prior. Whereas I2SB used
the degraded version (such as blurred, or masked) of the target image as the
originating boundary pair, we are free to design our own prior. Before we do
that, we outline below the conditions such a prior must satisfy to be an
effective candidate, based on the theoretical discussion in the preceding
sections:

  * •

Let $X_{0}$ be the target trjectory from $p_{\mathcal{A}}\left(X_{0}\right)$
and $X_{1}$ be the prior. Then, $X_{1}$ must be sampled from
$p_{\mathcal{B}}\left(X_{1}\mid X_{0}\right)$. In other words, the prior and
the target should be a pair \- there should be a clear dependence between the
two.

  * •

The prior should either be trivial to sample or be far cheaper to sample than
the target policy. In other words, the overhead of sampling
$p_{\mathcal{B}}\left(X_{1}\mid X_{0}\right)$ must be trivial compared to the
sampling time of $X_{0}$. (without violating the first condition)

  * •

The prior should be as close to the target distribution as possible. (without
violating the first and second conditions)

We are now ready to think about designing our priors. For the sake of
comparison, we can classify the prior distributions into the following kinds:

  1. 1.

Uninformative Prior: This has no structural information about the target
distribution, i.e., $p_{\mathcal{B}}\left(X_{1}\mid
X_{0}\right)=p_{\mathcal{B}}\left(X_{1}\right)$

    1. (a)

Random Prior: We use Gaussian white noise as our uninformative prior. This
should reduce our framework back to DDPM. We can use this as a control to
compare with the original Diffuser results.

  2. 2.

Informative Prior: This does share structural information with the target
trajectory and is closer to it than random white noise.

    1. (a)

Analytical Prior: A handcrafted or analytically calculated prior that is
trivial to sample.

    2. (b)

Learned Prior: Something that can be cheaply sampled by an inexpensive learned
model.

We also add the Number of Function Evaluations (NFE) as a variable separate
from the number of diffusion steps. Since SB posterior sampling can be done in
closed form (Equation 15), we test out smaller NFEs while keeping the same
number of diffusion steps. For DDPM, since sampling is iterative, the number
of function evaluations is always equal to the number of diffusion steps.

With this architecture and theoretical foundations in place, we can move on to
the experiments in the next section.

##  Chapter 3 Evaluation

###  1 Experimental Setup

####  1 Environment and Task

To evaluate our model against Diffuser’s default DDPM model, we design our
experiment to run on the Maze2D task in the D4RL [7] suite. The task involves
navigating a 2D maze by pushing a point mass to a goal location. We choose
this task for a few reasons:

  1. 1.

The original Diffuser paper primarily compares results on the Maze2D
benchmark.

  2. 2.

Since this is a maze navigation task, it is easy to construct and visualize
both handcrafted and learned model priors for this task, aiding us in figuring
out what’s happening.

  3. 3.

The task is one of goal-conditioned planning, which gets conveniently posed as
inpainting in the Diffuser architecture. Since the I2SB paper uses image
inpainting as one of the primary tasks, it makes for a much more equal
comparison, without the use of any guidance functions in Diffuser.

We use the normalized total score [7] obtained on the task over several
instantiations of the task to compare performance. We compare the performance
against the amount of data the two models are trained on to compare sample
efficiencies. The D4RL Maze2D environment generates trajectories to train the
samples on, and Diffuser uses the number of training steps to generate more or
less training trajectories. We, therefore, use training steps to represent the
amount of training data.

####  2 Task Details and Metrics

The environments are mazes of different sizes and complexity - open, umaze,
maze2d-medium, maze2d-large. They are visualized below in Figure 1.

The inputs of the task are the start location (green dot) and the goal
location (red dot). The output of the planning algorithm is a trajectory - a
set of state and action pairs. We sample in an open-loop fashion. The actions
returned by the planner are taken in the environment one by one by a waypoint
controller, and the environment returns a reward (which can be positive or
negative) for each step taken. The rewards are added up and divided by an
expert policy score, and the result is normalized to the 0 to 100 range. [7]

Figure 1: The four Maze2D environments in D4RL [7] \- open, umaze,
maze2d-medium, maze2d-large.

####  3 Prior Construction

We now describe the priors constructed for our experiments and also show their
trajectory on the maze2d-medium environment. In the images that follow, the
paths shown are the trajectories for the pair - the prior and the target - for
the same set of starting and ending points. The colors represent the
progression of the trajectory - points closer in color are closer in timestep.
The colorspaces for the prior and actual rollout are kept different for ease
of identification. For the priors, the start positions are blue and progress
in a gradient to the final position in red. For the target trajectories, we
use dark blue for start positions and progress in a gradient to light blue.

##### Analytical Prior

For the analytical prior, we handcraft a straight-line trajectory between the
start and goal locations. In other words, given a target trajectory $X_{1}$,
our prior is another trajectory $X_{0}$ that joins the ends of $X_{1}$ by a
straight line in the 2D maze environment, irrespective of the maze obstacles.
This is crafted by interpolating the 2D space in the maze to determine the
state space for each point, and using constant velocity in the direction of
the goal in the action space. The number of timesteps are kept the same as the
horizon.

This prior is trivial to evaluate, depends on the target trajectory, and has
structural information about the target - it lies in the same region of the
maze and connects the start and end points. It does not, however, have any
information about the structure of the maze environment itself.

In Figure 2 below, we show one sample of an analytical prior and its
corresponding target trajectory on maze2d-medium.

Figure 2: Left: Analytical prior trajectory $X_{0}$. Right: Target trajectory
$X_{1}$.

##### Learned Prior

For a learned model, we want something that is inexpensive to train and sample
and yet can approximate the target distribution to some degree. To that end,
we train a simple model with two linear layers separated by a ReLu layer,
trained to predict the prior trajectory in the same loss function and the same
trajectories as the diffusion model. The code for our simple prior model in
pyTorch is below - the observation dimension is the length of the states
vector, while the transition dimension is the length of the joint states and
actions vector:

⬇

mlp = nn.Sequential(

nn.Linear(2*observation_dim, horizon*transition_dim),

nn.LeakyReLU(0.1),

nn.Linear(horizon*transition_dim, horizon*transition_dim),

)

Figure 3 below shows a pair of learned prior and actual target rollout on
maze2d-medium.

Figure 3: Left: Learned prior trajectory $X_{0}$. Right: Target trajectory
$X_{1}$.

The approximate trajectories not only contain information about the start and
goal states, but also the structure of the maze itself, and are hence closer
to the target than the analytical prior above.

For demonstration, we chose a sample where the performance of the learned
model was good enough, just to show how much information it can contain.
Typically, the plans from such a simple model were much less precise and often
cut through obstacles.

##### Random Prior

Finally, for control, we also test using Gaussian white noise as a prior - the
same as what DDPM diffuses to. Each state and action value in the trajectory
of the random prior is simply a random number.

Figure 4: Left: Gaussian noise prior $X_{0}$. Right: Target trajectory
$X_{1}$.

###  2 Results

In the sections below, we discuss the results of benchmarking I2SB against
DDPM on planning tasks. We ran all benchmarks for 200 episodes each, and the
horizon was fixed at 256 for the maze2d-medium and 384 for the maze2d-large
environment. The scores are the normalized scores as defined in D4RL [7]
suite.

####  1 Uninformative Random Prior

##### Diffusing from Gaussian noise - Same NFE

Even though both of the models use the same networks and both start diffusing
from Gaussian noise, we see in Figure 5 a consistent drop in performance with
I2SB compared to DDPM. This suggests that when starting from pure Gaussian
noise, DDPM is a better formulation than I2SB in terms of sample quality, and
since the gap increases with training samples, also a better learner. We talk
about this later in Section 4.

Figure 5: Schrödinger Bridge in the I2SB formulation does not perform as well
as DDPM when starting from Gaussian prior.

##### Diffusing from Gaussian noise - across a range of NFEs

Next, we evaluate if SB can leverage the closed-form sampling (Equation 15) to
outperform DDPM at lower NFEs by keeping the number of diffusion steps larger.

From Figure 6 below, we see that DDPM struggles with NFE=1, and performance
does not improve with more training. With NFE=4 and above, DDPM can reach the
performance ceiling given enough data.

Figure 6: Score across different diffusion steps for DDPM.

For SB (Figure 7), we sample at varying NFEs while keeping the number of
diffusion steps (and therefore model complexity) constant. SB performs nearly
identically even at low NFEs. This suggests that closed-form sampling
(Equation 15) works well and prevents loss of performance at low NFEs.

Figure 7: Score across different NFEs for SB. Number of diffusion steps was
constant at 16.

####  2 Informative Priors

We now compare the priors we designed against each other, keeping NFE
constant. We see that at the lowest training steps, more informative priors
tend to perform better than less informative ones.

Figure 8: Learned priors perform better than analytical priors.

In Figure 9, we compare the performance of SB models directly against DDPM at
NFE=1. Even though DDPM beats I2SB at higher model capacities, here we see
that I2SB matches or surpasses the performance of DDPM on all three priors,
once again suggesting the strength of closed-form sampling.

Figure 9: I2SB can perform better than DDPM when the NFE is low.

###  3 Summary of Experimental Results

We started by asking if starting with a prior distribution instead of Gaussian
noise can help us with:

  1. 1.

Sample efficiency, i.e., can we get the same performance for less training
data?

  2. 2.

Sampling efficiency, i.e., can we get the same performance with faster
sampling times?

We can now begin to answer these questions.

####  1 Sample Efficiency

We first note from Figure 5 that given the same model complexity, the I2SB
formulation does no better than DDPM when starting from noise - in fact, it
does consistently worse across the number of training samples.

However, comparing the performance at NFE=1 from Figure 9, we can see that the
I2SB model does far better, especially at low training steps. Thus, at very
low NFEs, I2SB can maintain a much higher model complexity than DDBM, and thus
have better sample efficiency. This, of course, is a direct consequence of the
way I2SB does sampling through Equation 15, which allows I2SB to sample
multiple time steps in a single function evaluation/network call.

At the same time, as the allowance for NFE increases to 4 and above, we see
from Figure 6 and Figure 8 that DDPM beats I2SB even when I2SB uses a learned
prior distribution compared to DDPM starting from pure noise. This suggests
that DDPM is a better learner given enough model complexity.

####  2 Sampling Efficiency

Again, it is evident from Figure 6 and Figure 7 that I2SB has better sampling
efficiency only at very low NFEs. As soon as we go above 4 NFE, I2SB is not
able to catch up to the sample quality of DDPM, rendering the question of
sampling efficiency moot.

####  3 On Priors

We note from Figure 9 that the more informative priors do better as the number
of training steps goes up. We also note that out learned prior performs
slightly better than our handcrafted analytical prior in Figure 8.

####  4 On the Performance Gap between I2SB and DDPM

A key portion of our hypothesis was whether generation can be faster if
started from a closer distribution, with the implicit assumption that the
generating processes that do this transport are equally good, so that a closer
distribution means more efficient generation. We find that this implicit
assumption did not hold in the case of I2SB.

We noted in Figure 5 when comparing I2SB against DDPM, that when:

  1. 1.

The model complexity is kept the same for both, i.e., the same network size
and the same number of diffusion steps (and therefore NFE, since DDPM needs
one function evaluation per diffusion step.)

  2. 2.

And the model complexity is large enough, i.e., NFE is greater than just 1

  3. 3.

And both start from Gaussian white noise

DDPM outperforms I2SB irrespective of the training steps, with a consistent
gap. This suggests that DDPM as an algorithm is a better learner, at least
when allowed enough complexity in terms of diffusion steps.

We explain this in two ways:

  1. 1.

Since I2SB is an approximation of the general Schrödinger Bridge, it makes
some theoretical relaxations to permit tractability. For example, in Equation
13, the authors modify the reverse drift to be the Dirac Delta distribution
centered on data point $a$. They rely on the generalizing effects of neural
networks to counter this effect. We also note that Diffuser’s UNet, which we
used for our predictions, is a smaller UNet than I2SB’s original network,
because we work on much lower dimensional inputs.

DDPM on the other hand is a more exact model without such relaxations.

  2. 2.

Some recent works have pointed out a flaw in the sampling procedure of I2SB
and other similar direct diffusion models (DDB). [5], which introduces
Consistent Direct Diffusion Bridges that constantly guides the trajectory to
satisfy data consistency, note the following in their paper [5]:

> Regardless of the choice in constructing DDB, there is a crucial component
> that is missing from the framework. While the sampling process starts
> directly from the measurement (or equivalent), as the predictions
> $\hat{\boldsymbol{x}}_{0\mid t}=G_{\theta}\left(\boldsymbol{x}_{t}\right)$
> are imperfect and are never guaranteed to preserve the measurement
> condition, the trajectory can easily deviate from the desired path, while
> the residual blows up. Consequently, this may result in inferior sample
> quality, especially in terms of distortion.

The above two aspects of I2SB are also core to its formulation. While the
closed-form sampling and diffusing from a prior do lend it advantages, a
poorer learning process because of the approximations, and a flawed sampling
process leading to distortions over time, take away those advantages at higher
model complexity. We do show (Figure 9) however that despite being a poorer
learner, I2SB outperforms DDPM at very low NFE and moderate to high diffusion
steps.

##  Chapter 4 Discussion & Conclusion

In this thesis, we explored a novel way of planning from prior plans. We used
two existing frameworks - the Diffuser way of planning using diffusion models,
and the image-to-image bridged diffusion framework from I2SB, and added our
own interpretation of possible planning priors, to make the method work.
Though the particular formulation of Schrödinger Bridges we used did not work
as well, the overall method certainly has merit, and replacing I2SB with
bridging algorithms that are also more efficient learners might eventually
yield the results we sought initially.

Our findings may also be limited because we tried the method on one particular
task and in one particular environment, however, future work on higher
dimensional trajectory tasks might yield more interesting results.

We also did not evaluate the sampling time of the priors themselves but
instead assumed that operation to be trivially cheap. It could be an
interesting direction to evaluate how weaker but cheaper algorithms combined
with bridging methods could yield a stronger algorithm at a lower cost than
using diffusion alone.

We also did not go into the theoretical reasons for the performance of various
bridging methods under various conditions as it was beyond the scope of this
work, but it is a very active area of research right now.

Better formulations of bridging models will certainly provide us with a model
that efficiently generates plans from trivially learnable priors, which may
significantly speed up the planning process, especially in the case of closed-
loop scenarios. Aside form with readily available samples from a closer
distribution, such as time series forecasting, and modelling temporal data
such as video generation, could also benefit from such advances.

## References

  * [1] Michael S. Albergo, Nicholas M. Boffi and Eric Vanden-Eijnden  “Stochastic Interpolants: A Unifying Framework for Flows and Diffusions”  arXiv, 2023  DOI: 10.48550/ARXIV.2303.08797
  * [2] Valentin De Bortoli, James Thornton, Jeremy Heng and Arnaud Doucet  “Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling”, 2023  arXiv:2106.01357 [stat.ML]
  * [3] Tianrong Chen, Guan-Horng Liu and Evangelos A Theodorou  “Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory”  In _International Conference on Learning Representations_ , 2022
  * [4] Ching-An Cheng, Andrey Kolobov and Adith Swaminathan  “Heuristic-Guided Reinforcement Learning”  In _Advances in Neural Information Processing Systems_ 34 Curran Associates, Inc., 2021, pp. 13550–13563  URL: https://proceedings.neurips.cc/paper_files/paper/2021/file/70d31b87bd021441e5e6bf23eb84a306-Paper.pdf
  * [5] Hyungjin Chung, Jeongsol Kim and Jong Chul Ye  “Direct Diffusion Bridge using Data Consistency for Inverse Problems”  In _Advances in Neural Information Processing Systems_ , 2023
  * [6] Yilun Du et al.  “Learning Universal Policies via Text-Guided Video Generation”  In _arXiv e-prints_ , 2023, pp. arXiv–2302
  * [7] Justin Fu et al.  “D4RL: Datasets for Deep Data-Driven Reinforcement Learning”, 2020  arXiv:2004.07219 [cs.LG]
  * [8] Jonathan Ho, Ajay Jain and Pieter Abbeel  “Denoising Diffusion Probabilistic Models”  In _Advances in Neural Information Processing Systems_ 33 Curran Associates, Inc., 2020, pp. 6840–6851
  * [9] Michael Janner, Yilun Du, Joshua Tenenbaum and Sergey Levine  “Planning with Diffusion for Flexible Behavior Synthesis”  In _International Conference on Machine Learning_ , 2022
  * [10] Ivan Kapelyukh, Vitalis Vosylius and Edward Johns  “DALL-E-Bot: Introducing Web-Scale Diffusion Models to Robotics”  In _IEEE Robotics and Automation Letters_ 8.7 Institute of ElectricalElectronics Engineers (IEEE), 2023, pp. 3956–3963  DOI: 10.1109/lra.2023.3272516
  * [11] Guan-Horng Liu et al.  “I2SB: Image-to-Image Schrödinger Bridge”  In _arXiv preprint arXiv:2302.05872_ , 2023
  * [12] Gaurav Parmar et al.  “Zero-shot Image-to-Image Translation”, 2023  arXiv:2302.03027 [cs.CV]
  * [13] Karl Pertsch, Youngwoon Lee and Joseph J. Lim  “Accelerating Reinforcement Learning with Learned Skill Priors”  In _Conference on Robot Learning (CoRL)_ , 2020
  * [14] Aditya Ramesh et al.  “Zero-Shot Text-to-Image Generation”  In _Proceedings of the 38th International Conference on Machine Learning_ 139, Proceedings of Machine Learning Research  PMLR, 2021, pp. 8821–8831  URL: https://proceedings.mlr.press/v139/ramesh21a.html
  * [15] Chitwan Saharia et al.  “Palette: Image-to-Image Diffusion Models”  In _ACM SIGGRAPH 2022 Conference Proceedings_ , SIGGRAPH ’22  Vancouver, BC, Canada: Association for Computing Machinery, 2022  DOI: 10.1145/3528233.3530757
  * [16] E. Schrödinger  “Sur la théorie relativiste de l’électron et l’interprétation de la mécanique quantique”  In _Annales de l’institut Henri Poincaré_ 2.4 INSTITUT HENRI POINCARÉ ET LES PRESSES UNIVERSITAIRES DE FRANCE, 1932, pp. 269–310  URL: http://www.numdam.org/item/AIHP_1932__2_4_269_0/
  * [17] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan and Surya Ganguli  “Deep Unsupervised Learning using Nonequilibrium Thermodynamics”  In _Proceedings of the 32nd International Conference on Machine Learning_ 37, Proceedings of Machine Learning Research  Lille, France: PMLR, 2015, pp. 2256–2265
  * [18] Jiaming Song, Chenlin Meng and Stefano Ermon  “Denoising Diffusion Implicit Models”  In _International Conference on Learning Representations_ , 2021  URL: https://openreview.net/forum?id=St1giarCHLP
  * [19] Yang Song and Stefano Ermon  “Generative Modeling by Estimating Gradients of the Data Distribution”  In _Advances in Neural Information Processing Systems_ 32 Curran Associates, Inc., 2019
  * [20] Yang Song et al.  “Score-Based Generative Modeling through Stochastic Differential Equations”  In _International Conference on Learning Representations_ , 2021  URL: https://openreview.net/forum?id=PxTIG12RRHS
  * [21] Ikechukwu Uchendu et al.  “Jump-Start Reinforcement Learning”  In _Proceedings of the 40th International Conference on Machine Learning_ 202, Proceedings of Machine Learning Research  PMLR, 2023, pp. 34556–34583
  * [22] Tianhe Yu et al.  “Scaling Robot Learning with Semantically Imagined Experience”  In _arXiv preprint arXiv:2302.11550_ , 2023
  * [23] Linqi Zhou, Aaron Lou, Samar Khanna and Stefano Ermon  “Denoising Diffusion Bridge Models”, 2023  arXiv:2309.16948