# OptiRoute: A Heuristic-assisted Deep Reinforcement Learning Framework for
UAV-UGV Collaborative Route Planning

Md Safwan Mondal1, Subramanian Ramasamy1, Pranav Bhounsule1 *This work was
supported by ARO grant W911NF-14-S-003.1Md Safwan Mondal, Subramanian Ramasamy
and Pranav A. Bhounsule are with the Department of Mechanical and Industrial
Engineering, University of Illinois Chicago, IL, 60607 USA<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Unmanned aerial vehicles (UAVs) are capable of surveying expansive areas, but
their operational range is constrained by limited battery capacity. The
deployment of mobile recharging stations using unmanned ground vehicles (UGVs)
significantly extends the endurance and effectiveness of UAVs. However,
optimizing the routes of both UAVs and UGVs, known as the UAV-UGV cooperative
routing problem, poses substantial challenges, particularly with respect to
the selection of recharging locations. Here in this paper, we leverage
reinforcement learning (RL) for the purpose of identifying optimal recharging
locations while employing constraint programming to determine cooperative
routes for the UAV and UGV. Our proposed framework is then benchmarked against
a baseline solution that employs Genetic Algorithms (GA) to select rendezvous
points. Our findings reveal that RL surpasses GA in terms of reducing overall
mission time, minimizing UAV-UGV idle time, and mitigating energy consumption
for both the UAV and UGV. These results underscore the efficacy of
incorporating heuristics to assist RL, a method we refer to as heuristics-
assisted RL, in generating high-quality solutions for intricate routing
problems.

##  I Introduction

Unmanned aerial vehicles (UAVs) have rapidly evolved as an emerging
technology, finding profound applications in both military and civilian
sectors [1, 2, 3, 4]. They are critical for real-time sensing in scenarios
like traffic monitoring [5], border security [6], disaster management [7] and
forest fire surveillance [8], all of which demand continuous data
transmission. However, a major limitation of UAVs in such persistent
applications is their restricted operational time due to limited battery
capacity. This challenge can be mitigated by leveraging the synergies of
multi-agent systems that combine UAVs with unmanned ground vehicles (UGVs).
This collaborative approach can enhance the overall task efficacy and prolong
the UAVs’ operational longevity [9, 10].

In this work, we address a cooperative routing problem involving a
heterogeneous team of a UAV and a UGV to visit a set of designated task nodes
in the quickest possible time (see Fig. 1). The UAV operates under limited
battery life constraint and is supported by the UGV that acts as a mobile
recharging depot. The UGV is also restricted in terms of speed and can only
travel on the road network. The recharging process of UAV by UGV is not
instantaneous which adds an additional layer of complexity to the problem. The
challenge, therefore, is to meticulously strategize the routes for both the
entities, ensuring that the UAV can effectively get recharged from the UGV to
execute the mission optimally. This necessitates a comprehensive cooperative
routing framework to optimally plan UAV and UGV routes while synchronizing
their rendezvous recharging instances.

Figure 1: Illustration of a fuel constrained UAV-UGV cooperative routing
problem studied in this paper. DRL policy specifies rendezvous points, and
heuristic planners outline optimal UAV-UGV paths to cover all task points.

###  I-A Related works

Extensive research has been carried out on different variants of UAV-UGV
cooperative routing problem [11, 12, 13, 14]. Traditional methodologies, such
as Mixed Integer Linear Programming (MILP) [15], graph matching [16], and
multi-echelon formulation [12], have previously been employed to tackle this
type of combinatorial optimization problem (CO). In most studies, the
cooperative routing is modelled as a variant of vehicle routing problem what
makes it a NP-hard problem. These methods often fail to scale well with the
number of the task points and lacks practical applicability due to the
intricate combinatorial nature of the problem. Hence, in recent years,
learning based methodologies have gained popularity as a promising solution to
solve the combinatorial optimization problems.

Paul et al. [17, 18] utilized an encoder-decoder based graph neural network to
model mutlti-robot task allocation problem (MRTA) and showed learning based
algorithms can produce quality solutions in significantly less computational
time compare to the non-learning based baselines. Li et al. [19] employed a
deep reinforcement learning (DRL) method that leveraged attention mechanisms
to tackle the heterogeneous capacitated vehicle routing problem. Their
methodology outperformed traditional heuristic solutions in both quality and
computational efficiency. Wu et al. [20] investigated truck and drone based
last mile delivery problem using a reinforcement learning (RL) paradigm. They
split the optimization problem into customer clustering problem and routing
components and applied an encoder-decoder framework with RL to resolve it. Fan
et al. [21] employed a multi-head attention mechanism coupled with a DRL
policy to design routes for an energy-limited UAV, however they assumed a
fixed set of recharging locations for the UAV.

In our prior research [22, 23], we introduced a two-tiered optimization
structure that streamlined the collaborative routing challenge. This framework
determined the UGV route in the primary layer and the UAV route in the
secondary layer using heuristic methods. We further demonstrated that
optimizing the recharging instances between the UAV and UGV can enhance the
cooperative routing solution’s quality [24]. In this paper, we have coupled
heuristic strategy-based route planners with a learning framework that has
produced an optimal strategy for UAV-UGV recharging rendezvous planning to
achieve the best cooperative route. Given the complexity of the task, we have
harnessed the strengths of both heuristic and learning methods to navigate the
scenario and have compared our solution with a non-learning baseline, which
has shown substantial improvement in solution quality. To this end, we present
the following novel contributions:
1\. We have formulated the fuel constrained UAV-UGV cooperative routing
problem as a Markov Decision Process (MDP) so that the recharging rendezvous
problem can be learned using a policy gradient RL approach. To the best of our
knowledge, this is one of the first work to address fuel constrained UAV-UGV
cooperative routing problem with RL.
2\. We have adopted an encoder-decoder based transformer architecture within
the learning policy network that can help in more efficient data extraction
from the scenario to select the appropriate refuel stop locations.
3\. We have integrated the advantages of heuristics with RL framework by
modelling the energy-constrained vehicle routing problem and solving it
through constraint programming heuristics to get the UAV-UGV routes based on
the recharging stop locations chosen by the RL policy.
4\. The comparison with a non-learning baseline method shows the effectiveness
and reliability of the proposed algorithm.

##  II Problem Formulation

###  II-A Problem overview

Given a set of task points $\mathcal{M}=\\{m_{0},m_{1},...,m_{n}\\}$ situated
on a road network, the objective is to strategize a collaborative operation
between a UAV $A$ and a UGV $G$ to optimally visit all the task points in the
scenario. The vehicles have heterogeneous characteristics as the UAV boasts
higher speed $v^{a}$ but operates within a restricted battery capacity
$F^{a}$. It can, however, recharge its battery from the UGV, which travels at
a slower pace $v^{g}$ exclusively on the road network. For recharging, the UAV
docks onto the UGV’s landing platform and once rejuvenated, it resumes its
flight by taking off from UGV. The duration of this recharging is contingent
upon the UAV’s battery level prior to docking. The UAV’s energy consumption
profile,
$\mathcal{P}^{a}=0.0461{v^{a}}^{3}-0.5834{v^{a}}^{2}-1.8761v^{a}+229.6$ being
a inversely proportional of its velocity $v^{a}$ adds another layer of
complexity to the problem as the hovering consumes more energy than
transiting, hence discouraged during the task. Considering these multifaceted
challenges, which largely revolve around the recharging interplay between the
UAV and UGV, our objective is to develop an efficient recharging strategy (
RendezvousPlanner ) which determines where and when the UAV can rendezvous
with UGV for recharging. Based on this rendezvous, routes for both vehicles
$\tau^{a}\in A,\tau^{g}\in G$ are constructed by the UAVPlanner and UGVPlanner
to ensure comprehensive coverage of all task points in the shortest time span,
while conserving energy and minimizing idle periods ( UAV hovering, UGV
waiting ) during recharging sessions.

###  II-B MDP formulation

Figure 2: MDP for heuristics-assisted UAV-UGV cooperative routing

The rendezvous planning of this cooperative routing can be formulated as a
MDP, since the refueling stop locations are chosen sequentially and its
selection influences the state of entire task scenario (see Fig. 2) The
components of the MDP can be defined as a tuple
$<\mathcal{\mathcal{missing}}{S},\mathcal{A},\mathcal{R},\mathcal{T}>$, as
following:
1) State Space $(\mathcal{S})$: In our MDP at any decision making step, the
state of the environment $s_{t}\in\mathcal{S}$ is defined as,
$s_{t}=(\\{x^{a}_{t},y^{a}_{t}\\}\in A,\\{x^{g}_{t},y^{g}_{t}\\}\in
G,\\{x^{i}_{t},y^{i}_{t},d^{i}_{t}\\})\in\mathcal{M})$, where,
$\\{x^{a}_{t},y^{a}_{t}\\}$ and $\\{x^{g}_{t},y^{g}_{t}\\}$ represent the
current coordinates of the UAV and UGV respectively and
$\\{x^{i}_{t},y^{i}_{t},d^{i}_{t}\\}$ highlights the coordinates of task
points and their visitation status, $d^{i}_{t}$ will be 1 if a task point is
already visited or 0 if not.
2) Action Space $(\mathcal{A})$: Selection of a task node as the refuel stop
$m^{r}$ is considered to be the action $a_{t}\in\mathcal{A}$. Specifically,
$a_{t}$ is defined as the index of the task node that is selected as
rendezvous location at that decision making step, i.e, $a_{t}=i,\ \forall\
m^{r}_{i}={m_{i}}\in\mathcal{M}$.
3) Reward ($\mathcal{R}$): By keeping congruity with the objective of the
cooperative routing problem to reduce total task completion time, we set the
reward $r_{t}\in\mathcal{R}$ as the negative value of UAV route time
$t^{a}_{route}$. This UAV route time is obtained by solving a vehicle routing
problem (VRP) problem based on the rendezvous location $m^{r}_{t}$ chosen in
the action $a_{t}$. The recharging time $t^{a}_{r}$ is calculated based on
UAV’s fuel consumption and subtracted from the reward function. Our reward
function also deducts $t^{a}_{idle}$ and $t^{g}_{idle}$ to discourage UAV’s
hovering and UGV’s waiting time period during rendezvous. For effective reward
shaping, we provide a significant positive bonus, $D$, upon task completion
and impose a substantial penalty, $P$, for task failure. So, the reward
$r^{t}$ can be written be as:
$r_{t}=r(s_{t},a_{t})=-t^{a}_{route}-t^{a}_{r}-t^{a}_{idle}-t^{g}_{idle}+D-P$.
Here, when the task is successfully completed, $D=10000,P=0$. In the case of
task failure, $D=0,P=1000$. This reward structure promotes successful task
completion and aids in achieving faster learning convergence.
4) Transition $(\mathcal{T})$: The next state $s_{t+1}$ in the transition
function $\mathcal{T}(s_{t+1}|s_{t},a_{t})$ depends on the UAV route
$\tau^{a}_{t}$ and UGV route $\tau^{g}_{t}$ obtained from route planners
UAVPlanner and UGVPlanner based on selected action $a_{t}$. This is where we
integrate routing heuristics (see section IV) with RL framework for solving
the cooperative route. In transition, the position of UAV and UGV is updated
and the visitation state of the mission points is updated based on the UAV and
UGV routes as follow: $s_{t+1}=(\\{x^{a}_{t+1},y^{a}_{t+1}\\}\in
A,\\{x^{g}_{t+1},y^{g}_{t+1}\\}\in
G,\\{x^{i}_{t+1},y^{i}_{t+1},d^{i}_{t+1}\\})\in\mathcal{M})$, here
$d^{i}_{t+1}=1,\ \text{ for unvisited mission points}\
m_{i}\in\tau^{a}_{t}|\tau^{g}_{t}$ and other elements is 0. Since
stochasticity is not assumed the transition is considered to be deterministic.

Figure 3: Architecture of the proposed policy network. Encoder has one self
attention layer to generate input scenario embedding, while decoder contains
one self attention to yield context and one cross attention layer to generate
attention score.

##  III Reinforcement learning framework

In our study, we have implemented encoder-decoder based transformer
architecture with RL algorithm to parameterize and learn the optimal policy
$\pi_{\theta}$ with a trainable parameter $\theta$. Starting from the initial
state $s_{0}$, the policy $\pi_{\theta}=Pr(a_{t}|s_{t},\theta_{t}=\theta)$
takes action $a_{t}$ to select appropriate refuel stops based on the scenario
state $s_{t}$ until the terminal state $s_{\tau}$ is reached. The final
solution of the policy network consists of a series of refuel stops chosen
sequentially which can be represented as joint probability distribution as
follow:

$Pr(s_{\tau}|s_{0})=\prod\limits_{t=0}^{\tau-1}\pi_{\theta}(a_{t}|s_{t})Pr(s_{t+1}|s_{t},a_{t})$
(1)

###  III-A Encoder-Decoder Transformer architecture

Our policy network is consists of two main components: an encoder and a
decoder. The encoder translates the task point coordinates into a more
nuanced, high-dimensional embedding for better feature extraction. On the
other hand, the decoder in the policy network $\pi_{\theta}$, discerns
suitable rendezvous locations based on the encoder embedding and the
contextual information extracted from the scenario’s current state. A detailed
description of the policy architecture (see Fig. 3) is explained here:

####  III-A1 Encoder

We have utilized a multi head attention (MHA) mechanism [25, 26] in the
encoder for higher dimensional representation of the raw features of the
problem instance. The encoder takes the normalized 2D coordinates of the task
points as an input $X=(o_{i}=\\{x_{i},y_{i}\\},\forall\ m_{i}\in\mathcal{M})$
and linearly projects them to a higher dimension space with dimension
$d_{h}=32$ to create input embedding $h^{0}_{i}$. This input embedding is
subsequently transformed using a single-layer multi-head self-attention
mechanism. This transformation process yields $h^{L}_{i}$, allowing a more
detailed understanding of the relationships among the task points. The
Attention layer calculates three vectors, query, key and value from the input
node embedding $h^{0}_{i}$ with the dimension of query/key, value as,
$d_{q}/d_{k}=d_{v}=\frac{d_{h}}{M}$, here $M=4$, number of attention heads.
For each head $j\in{1,2,...M}$, the attention scores $Z_{j}$ are calculated
between query $q_{i,j}$ and key $k_{i,j}$ and concatenated together to the
original embedding dimension $h^{0}_{i}$. The calculations are shown here:

$\displaystyle q_{i,j}=h^{0}_{i}W^{q}_{j},\ k_{i,j}=h^{0}_{i}W^{k}_{j},\
v_{i,j}=h^{0}_{i}W^{v}_{j}$ (2) $\displaystyle
Z_{j}=\text{softmax}\left(\frac{q_{i,j}{k_{i,j}}^{T}}{\sqrt{d_{k}}}\right)v_{i,j}$
(3) $\displaystyle
h^{L}_{i}=MHA(h^{0}_{i})=\text{Concat}(Z_{1},Z_{2},...,Z_{j})$ (4)

Here, $q_{i,j},k_{i,j}\ \text{and}\ v_{i,j}$ are the query, key and value
respectively in head $j$ and $W^{q}_{j}\in\mathbb{R}^{d_{h}\times
d_{q}},W^{k}_{j}\in\mathbb{R}^{d_{h}\times d_{k}}$ and
$W^{v}_{j}\in\mathbb{R}^{d_{h}\times d_{v}}$ are the trainable parameter
matrices in the attention layer. The self attention output $h^{L}_{i}$ has a
residual skip connection around them, followed by Layer-Normalization (LN).
Then for the richer interpretation of feature space, the output $h^{r}_{i}$ is
refined through a Feed Forward (FF) layer followed by another residual skip
connection and Layer-Normalization that yields the encoder embedding
$h^{f}_{i}$ of the problem instance which is leveraged in the decoder.

$\displaystyle h^{r}_{i}=LN(h^{0}_{i}+MHA(h^{0}_{i})),$ (5) $\displaystyle
h^{f}_{i}=LN(h^{r}_{i}+FF(\text{ReLU}(h^{r}_{i})))$ (6)

####  III-A2 Decoder:

During each decision-making step, the decoder determines the probability of
selecting each available node as an action based on the encoder’s node
embedding $h^{f}_{i}$ and a context vector, which provides insights into the
current scenario state. The decoder employs multi head self attention layer to
extract state features as the context $H^{con}_{t}$ which is achieved by
linear transformation of current state $s_{t}$ into $H^{0}_{i}$ and subsequent
processing through multi head self attention layer followed by skip residual
connection and Layer-Normalization as shown in following equations:

$\displaystyle H^{0}_{i}=\mbox{linear}(s_{t}),\
H^{con}_{t}=LN(H^{0}_{i}+MHA(H^{0}_{i}))$ (7)

Then this context vector $H^{con}_{t}$ is treated as the Query and the encoder
node embedding $h^{f}_{i}$ is considered as the Key/Value pair for a multi-
head heterogeneous cross attention layer to generate $H^{attn}_{t}$.

$\displaystyle H^{attn}_{t}=MHA(H^{con}_{t}W^{Q}_{j},\ h^{f}_{i}W^{K}_{j},\
h^{f}_{i}W^{V}_{j})$ (8)

Following the original transformation network [26, 25], the cross attention
score $H^{attn}_{t}$ has a skip connection and Layer-Normalization. This is
succeeded by a Feed Forward (FF) layer and another Layer-Normalization,
further refining it to $H^{f}_{t}$. At the final step, the attention score is
linearly projected to the action space dimension and is then averaged across
the nodes.

$\displaystyle H^{r}_{t}=LN(H^{attn}_{t}+H^{con}_{t})$ (9) $\displaystyle
H^{f}_{t}=LN(H^{r}_{t}+FF(\text{ReLU}(H^{r}_{t})))$ (10) $\displaystyle
H^{out}_{t}=\text{mean}(\text{linear}(H^{f}_{t}))$ (11)

An important step in this decoding step is to mask the invalid action nodes by
setting its attention score to negative infinity. In our case, any task point
that is out of reach of UAV and UGV from its current position and fuel
capacity is considered as invalid at that action step. Finally, a softmax
layer outputs the probability distribution over the action nodes,
$Pr(a_{t}|s_{t},\theta_{t}=\theta)=\text{softmax}(H^{out}_{t})$, indicating
the likelihood of each node to be selected as the refueling point. We have
utilized sampling strategy to sample the node as action according to their
probability in the decoding stage.

###  III-B Training method

The training of our policy network leverages the REINFORCE policy gradient
method, a seminal approach in the realm of reinforcement learning. At its
core, REINFORCE seeks to optimize policies by directly adjusting policy
parameters in the direction that improves expected rewards. The training
process ( see Algorithm 1 ) begins by initializing the policy parameters,
often randomly. We execute an episode under the current policy, producing
state, action, and reward sequences. Each step’s return is calculated, and the
policy parameters are adjusted based on the gradient of the log probability of
actions, scaled by these returns. This iterative process makes actions with
higher returns more likely in future episodes. The method’s strength lies in
its model-free nature, eliminating the need for explicit knowledge of the
environment’s dynamics. However, in our study future enhancements may include
baseline networks with algorithms like Proximal Policy Optimization (PPO) or
Actor-Critic method for better convergence.

Algorithm 1 Policy network training using REINFORCE

1:Policy network $\pi_{\theta}$, epochs $E$, batches $B$, episode length $T$,
learning rate $\alpha$, discount factor $\gamma$

2:Trained policy network $\pi_{\theta^{{}^{\prime}}}$

3:Initialize policy $\pi_{\theta}$ parameters $\theta$

4:for epoch in $1\dots E$ do

5: Initialize trajectory set $\mathcal{E}=[\ ]$

6: for instance in $1\dots B$ do

7: Initialize $s_{0}$ and $t\leftarrow 0$

8: Initialize trajectory $\tau=[\ ]$

9: while $t<T$ do

10: Get action $a_{t}\sim\pi_{\theta}(a_{t}|s_{t})$

11: Obtain $r_{t}$ and $s_{t+1}$

12: Store $(s_{t},a_{t},r_{t})$ in $\tau$, $t=t+1$

13: end while

14: Append $\tau$ to $\mathcal{E}$

15: end for

16: Compute gradient:
$\nabla_{\theta}J=\frac{1}{B}\sum_{\tau\in\mathcal{E}}\sum_{(s,a,r)\in\tau}\nabla_{\theta}\log\pi_{\theta}(a|s)\sum_{t^{\prime}=t}^{T}\gamma^{t^{\prime}-t}r_{t^{\prime}}$

17: Update policy parameters: $\theta\leftarrow\theta+\alpha\nabla_{\theta}J$

18:end for

##  IV Routing Heuristics

Based on the refuel stops selected by the policy network (RendezvousPlanner),
we can model vehicle routing problem with constraints to obtain the UGV & UAV
routes which in turn give us the reward of our MDP in the learning framework.
Since the UGV is slower than the UAV, we followed the ‘UGV first, UAV second’
heuristics approach of cooperative routing to construct the UGV route as the
initial step and then construct the UAV route based on the spatio-temporal
components of UGV route.

####  IV-1 UGV route planner

The UGVPlanner generates the UGV path by connecting the UGV’s current position
to the chosen refuel stops. Since the UGV is confined to the road network, the
waypoints between its present position and the refuel stop provide its route’s
spatial aspect. Given the UGV’s consistent speed, its temporal journey can
also be determined. To align with the UAV for refueling, the UGV may wait at
the refuel station if necessary, however waiting at the refuel stops is
discouraged to reduce idle period. The arrival time of the UGV at the refuel
stop is fed to UAVPlanner which models the UAV routing as the energy
constrained vehicle routing problem with time windows (E-VRPTW).

####  IV-2 UAV route planner

The E-VRPTW formulation can be depicted using graph theory. Here, the task
points act as the vertices $V=\\{S,0,1,...D\\}$ and $E=\\{(i,j)\,\|\
i,j\,\in\,V,i\neq j\\}$ denotes the edges connecting vertices $i$ and $j$. We
assign non-negative arc cost between vertices $i$ and $j$ as $t_{ij}$
(traversal time) and the decision variable as $x_{ij}$ that indicates whether
a vehicle transits from $i$ to $j$. The UAV commences it’s journey at starting
point $S$, visits the task points and when needed it terminates its route to
recharge from the UGV at refuel stop $D$, which is bounded by a time-window
due to UGV’s slower pace. The objective (Eq. 12) is minimizing the cumulative
travel duration while dropping least number of task points by penalizing each
task point drop with a large penalty $P$; here $y_{i}$ is a binary variable
that’s set to 1 if a point is visited or 0 otherwise. We’ve established energy
constraint in Eqs. 13 \- 14 to make sure that the UAV never runs out of its
fuel and its fuel consumption follows the UAV’s power consumption profile
during traversal (Eq. 15). The time-window condition in Eq. 16 makes the UAV
visit the UGV only after its arrival at refuel stop $D$. Eq. 17 says that the
cumulative arrival time at $j^{th}$ node is equal to the sum of cumulative
time at the node $i$, $t_{i}$ and the travel time between them $t_{ij}$. In
both Eq. 15 & Eq. 17 we have applied Miller-Tucker Zemlin (MTZ) formulation
[27] by adding large constant $L_{1},L_{2}$ for sub-tour elimination in the
UAV route. The other generic constraints for the UAV route like flow
conservation are detailed in our previous work [28].

Objective:

$\min\sum_{i}\sum_{j}t_{ij}x_{ij}+P\sum_{i}(1-y_{i})\quad\forall i,j\in V$
(12)

Energy constraints:

$f^{a}_{j}=F^{a},\quad j\in D$ (13) $0\leq f^{a}_{j}\leq F^{a},\quad\forall
j\in V\setminus\\{S,D\\}$ (14) $\displaystyle f^{a}_{j}$ $\displaystyle\leq
f^{a}_{i}-\left(\mathcal{P}^{a}(v^{a})t_{ij}x_{ij}\right)$
$\displaystyle+L_{1}\left(1-x_{ij}\right),\quad\forall i\in V,j\in
V\setminus\\{S,D\\}$ (15)

Time window constraints:

$t_{j,start}\leq t_{j}\leq t_{j,end}\ ,\quad\forall j\in D$ (16) $t_{j}\geq
t_{i}+\left(t_{ij}x_{ij}\right)-L_{2}\left(1-x_{ij}\right),\quad\forall i\in
V,j\in V$ (17)

The UAV route is calculated by the UAVPlanner by solving the above E-VRPTW
formulation with Google OR-Tools ™ CP-SAT solver [29] that uses constrained
programming (CP) and we leverage TABU SEARCH heuristics in the solver for
avoiding the local optimal solution.

Distinct from earlier cooperative routing studies [9, 12], our approach
leverages mobile recharging, turning recharging duration into opportunities
for mission point visitation. Once the UAV-UGV route is determined, the UAV’s
refueling time is computed based on its in-flight fuel consumption. This
refueling time is projected as a distance along road network where UGV
transports the UAV, effectively recharging it while concurrently visiting task
points along the road. The end location of the recharging process will be the
next take off point of the UAV and the route time, recharging time, UAV
hovering time and UGV waiting time are fed back to the MDP reward function and
we will move forward to the next cycle in MDP.

##  V Simulation Results

In this section, we simulate the fuel-constrained cooperative routing of a
UAV-UGV team across a 20 km $\times$ 20 km real world site. Starting from a
depot, the team aims to visit task points shown as black dots in Fig. 5. The
UAV and UGV moves at $v^{a}=10$ m/s and $v^{g}=4.5$ m/s speed respectively and
UAV has a fuel capacity of 287.7 kJ. The recharging instances between UAV and
UGV are determined by RendezvousPlanner and their route is determined by
UAVPlanner and UGVPlanner respectively as previously discussed. The goal is to
accomplish the task in minimum possible time while ensuring that the UAV and
UGV spend minimal idle time hovering or waiting during their recharging
rendezvous instances. We evaluate the following metrics: 1) the total task
completion time as the main metric 2) the idle time spent percentage (idle
time / task completion time) by UAV and UGV during rendezvous 3) the UAV and
UGV energy consumption metrics during the entire task period. Given the unique
complexity and specificity of our challenge, there are no standard benchmarks
available, nor could an exact solution be discerned. Hence we compare our RL
based results with results derived from a Genetic algorithm (GA) based non-
learning baseline. The details about the implementation of GA for this
cooperative routing can be found in our previous study [24]. All computational
procedures were executed using a Python simulation on a 3.7 GHz Intel Core i9
processor, equipped with 32 GB RAM, operating on a 64-bit system.

The training of RL is conducted in 500 episodes with different learning rates
$lr=0.001,0.005,0.01$. The learning rate significantly influences the RL
algorithm’s convergence speed. As illustrated in Fig. 4, a lower $lr=0.001$
results in a prolonged convergence, whereas an aggressive $lr=0.01$ converges
swiftly but tends to overfit, compromising policy exploration. The
intermediate $lr=0.005$ strikes a balance between exploration and
exploitation, emerging as the most effective for deriving an optimal policy.
Consequently, we adopt the $\pi_{\theta}$ policy trained with $lr=0.005$ for
coordinating the UAV-UGV rendezvous in our task scenario.

Figure 4: Mean episode reward in training process  TABLE I: Comparison of Routing metrics  Route Metrics |  DRL method |  GA method
---|---|---
Total task completion time (min) |  143.00 |  186.00
Mission points dropped |  0 |  0
UAV travel time (min) |  84.72 |  99
UAV hovering time (min) |  2.41 |  4
UAV hovering energy cost (kJ) |  33.26 |  55.20
UAV total energy cost (kJ) |  1039.73 |  1231.32
UGV travel time (min) |  71.92 |  52.00
UGV waiting time (min) |  15.22 |  61.00
UGV waiting energy cost (kJ) |  325.30 |  1304.06
UGV total energy cost (kJ) |  19347.13 |  19911.56

Figure 5: Cooperative routes of UAV-UGV at recharging instances A) DRL method
B) GA method. The route animation can be found at http://tiny.cc/x55bvz

Although, both the GA and DRL frameworks are able to come up with routing
strategy to cover entire task scenario without drooping any task nodes, Fig. 5
shows the difference of the cooperative routes obtained by them. DRL optimally
determines the recharging rendezvous between the UAV and UGV, ensuring even-
handed utilization of both in the routing process. In contrast to the GA
approach, DRL’s chosen refuel stops are further from previous rendezvous
points. Consequently, the UGV travels longer to rendezvous locations hence
doesn’t endure prolonged wait times at refuel stops for UAV to land. From
table II reveals that with strategic rendezvous planning, DRL requires only 4
recharging instances to complete the task, while GA demands 5. More recharging
events add extra recharging and UAV detour times, leading to a prolonged
overall task completion time (see table I). Also, owing to the less detours
and less hovering UAV total energy consumption is 15% lesser in DRL method.

TABLE II: Comparison of Rendezvous metrics Rendezvous Metrics |  DRL method |  GA method
---|---|---
No. of recharge stops |  4 |  5
Recharge time (min) |  55.87 |  73.00
UAV idle time (%) |  1.68 |  2.15
UGV idle time (%) |  10.64 |  32.79

The impact of the rendezvous policy can be clearly understood in table II. In
DRL rendezvous policy, UAV and UGV spend less amount of time in remain idle
either in hovering or waiting. Particularly, the UGV idle time is 3 times less
in the DRL method what fulfills the objective of this work. Hence we can say,
DRL method is able to learn the given environment to figure out the optimal
recharging policy between the UAV and UGV which ultimately impacts the overall
the cooperative route positively by reducing the task completion time and
vehicle idleness and energy consumption.

##  VI Conclusion & Future work

In this study, we introduced a novel framework that combined heuristic
guidance with deep reinforcement learning to address the fuel-constrained UAV-
UGV cooperative routing problem. We demonstrated the effectiveness of using an
encoder-decoder transformer architecture for making rendezvous decisions,
while a heuristic planner efficiently devised routes for the UAV and UGV to
minimize total mission time. By harnessing the decision-making capabilities of
reinforcement learning and the route generation through heuristics, our
approach aptly navigated the complexities of such a problem. With our DRL-
based framework, we achieved a 23% reduction in task completion time compared
to the non-learning GA baseline. Additionally, through strategic rendezvous
planning, we significantly minimized the UAV-UGV idle periods, leading to a
considerable decrease in their energy consumption. In future, we plan to
incorporate batch processing in RL to enhance training, aiming for a more
robust solution. This will allow us to make more broader claim about the
viability of our approach. Furthermore, we intend to adapt our framework for
dynamic routing to account for unpredictable variations.

## References

  * [1] Yao Liu, Zhihao Luo, Zhong Liu, Jianmai Shi, and Guangquan Cheng.  Cooperative routing problem for ground vehicle and unmanned aerial vehicle: The application on intelligence, surveillance, and reconnaissance missions.  IEEE Access, 7:63504–63518, 2019.
  * [2] Daniel H Stolfi, Matthias R Brust, Grégoire Danoy, and Pascal Bouvry.  Uav-ugv-umv multi-swarms for cooperative surveillance.  Frontiers in Robotics and AI, 8:616950, 2021.
  * [3] Pratap Tokekar, Joshua Vander Hook, David Mulla, and Volkan Isler.  Sensor planning for a symbiotic uav and ugv system for precision agriculture.  IEEE transactions on robotics, 32(6):1498–1511, 2016.
  * [4] Yu Wu, Shaobo Wu, and Xinting Hu.  Cooperative path planning of uavs & ugvs for a persistent surveillance task in urban environments.  IEEE Internet of Things Journal, 8(6):4906–4919, 2020.
  * [5] Anuj Puri, KP Valavanis, and M Kontitsis.  Statistical profile generation for traffic monitoring using real-time uav based video data.  In 2007 Mediterranean Conference on Control & Automation, pages 1–6. IEEE, 2007.
  * [6] Omer Ozkan and Muhammed Kaya.  Uav routing with genetic algorithm based matheuristic for border security missions.  An International Journal of Optimization and Control: Theories & Applications (IJOCTA), 11(2):128–138, 2021.
  * [7] Milan Erdelj, Enrico Natalizio, Kaushik R Chowdhury, and Ian F Akyildiz.  Help from the sky: Leveraging uavs for disaster management.  IEEE Pervasive Computing, 16(1):24–32, 2017.
  * [8] Chi Yuan, Youmin Zhang, and Zhixiang Liu.  A survey on technologies for automatic forest fire monitoring, detection, and fighting using unmanned aerial vehicles and remote sensing techniques.  Canadian journal of forest research, 45(7):783–792, 2015.
  * [9] Parikshit Maini and PB Sujit.  On cooperation between a fuel constrained uav and a refueling ugv for large scale mapping applications.  In 2015 international conference on unmanned aircraft systems (ICUAS), pages 1370–1377. IEEE, 2015.
  * [10] Mingjia Zhang, Huawei Liang, and PengFei Zhou.  Cooperative route planning for fuel-constrained ugv-uav exploration.  In 2022 IEEE International Conference on Unmanned Systems (ICUS), pages 1047–1052. IEEE, 2022.
  * [11] Jianqiang Li, Genqiang Deng, Chengwen Luo, Qiuzhen Lin, Qiao Yan, and Zhong Ming.  A hybrid path planning method in unmanned air/ground vehicle (uav/ugv) cooperative systems.  IEEE Transactions on Vehicular Technology, 65(12):9585–9596, 2016\.
  * [12] Yao Liu, Zhong Liu, Jianmai Shi, Guohua Wu, and Witold Pedrycz.  Two-echelon routing problem for parcel delivery by cooperated truck and drone.  IEEE Transactions on Systems, Man, and Cybernetics: Systems, 51(12):7450–7465, 2020.
  * [13] Satyanarayana G Manyam, Kaarthik Sundar, and David W Casbeer.  Cooperative routing for an air–ground vehicle team—exact algorithm, transformation method, and heuristics.  IEEE Transactions on Automation Science and Engineering, 17(1):537–547, 2019.
  * [14] Subramanian Ramasamy, Jean-Paul F Reddinger, James M Dotterweich, Marshal A Childers, and Pranav A Bhounsule.  Coordinated route planning of multiple fuel-constrained unmanned aerial systems with recharging on an unmanned ground vehicle for mission coverage.  Journal of Intelligent & Robotic Systems, 106(1):30, 2022.
  * [15] Kaarthik Sundar, Saravanan Venkatachalam, and Sivakumar Rathinam.  Formulations and algorithms for the multiple depot, fuel-constrained, multiple vehicle routing problem.  In 2016 American Control Conference (ACC), pages 6489–6494. IEEE, 2016.
  * [16] Payam Ghassemi and Souma Chowdhury.  Multi-robot task allocation in disaster response: Addressing dynamic tasks with deadlines and robots with range and payload constraints.  Robotics and Autonomous Systems, 147:103905, 2022.
  * [17] Steve Paul and Souma Chowdhury.  A scalable graph learning approach to capacitated vehicle routing problem using capsule networks and attention mechanism.  In International Design Engineering Technical Conferences and Computers and Information in Engineering Conference, volume 86236, page V03BT03A045. American Society of Mechanical Engineers, 2022.
  * [18] Steve Paul, Wenyuan Li, Brian Smyth, Yuzhou Chen, Yulia Gel, and Souma Chowdhury.  Efficient planning of multi-robot collective transport using graph reinforcement learning with higher order topological abstraction.  arXiv preprint arXiv:2303.08933, 2023.
  * [19] Jingwen Li, Yining Ma, Ruize Gao, Zhiguang Cao, Andrew Lim, Wen Song, and Jie Zhang.  Deep reinforcement learning for solving the heterogeneous capacitated vehicle routing problem.  IEEE Transactions on Cybernetics, 52(12):13572–13585, 2021.
  * [20] Guohua Wu, Mingfeng Fan, Jianmai Shi, and Yanghe Feng.  Reinforcement learning based truck-and-drone coordinated delivery.  IEEE Transactions on Artificial Intelligence, 2021.
  * [21] Mingfeng Fan, Yaoxin Wu, Tianjun Liao, Zhiguang Cao, Hongliang Guo, Guillaume Sartoretti, and Guohua Wu.  Deep reinforcement learning for uav routing in the presence of multiple charging stations.  IEEE Transactions on Vehicular Technology, 2022.
  * [22] Subramanian Ramasamy, Md Safwan Mondal, Jean-Paul F Reddinger, James M Dotterweich, James D Humann, Marshal A Childers, and Pranav A Bhounsule.  Solving vehicle routing problem for unmanned heterogeneous vehicle systems using asynchronous multi-agent architecture (a-teams).  In 2023 International Conference on Unmanned Aircraft Systems (ICUAS), pages 95–102. IEEE, 2023.
  * [23] Md Safwan Mondal, Subramanian Ramasamy, and Pranav Bhounsule.  A bilevel optimization framework for fuel-constrained uav-ugv cooperative routing: Planning and experimental validation.  arXiv preprint arXiv:2303.02315, 2023.
  * [24] Subramanian Ramasamy, Md Safwan Mondal, Jean-Paul F Reddinger, James M Dotterweich, James D Humann, Marshal A Childers, and Pranav A Bhounsule.  Heterogenous vehicle routing: comparing parameter tuning using genetic algorithm and bayesian optimization.  In 2022 International Conference on Unmanned Aircraft Systems (ICUAS), pages 104–113. IEEE, 2022.
  * [25] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.  Neural machine translation by jointly learning to align and translate.  arXiv preprint arXiv:1409.0473, 2014.
  * [26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [27] Clair E Miller, Albert W Tucker, and Richard A Zemlin.  Integer programming formulation of traveling salesman problems.  Journal of the ACM (JACM), 7(4):326–329, 1960.
  * [28] Md Safwan Mondal, Subramanian Ramasamy, James D. Humann, Jean-Paul F. Reddinger, James M. Dotterweich, Marshal A. Childers, and Pranav A. Bhounsule.  Cooperative multi-agent planning framework for fuel constrained uav-ugv routing problem, 2023.
  * [29] Google.  Google OR-tools.  https://developers.google.com/optimization, 2021.  Online; accessed Feb 2, 2021.


# Identified hadron production at mid-rapidity in Au+Au collisions at
$\sqrt{s_{NN}}=54.4$ GeV at STAR

Krishan Gopal (for the STAR Collaboration)  Department of Physics, Indian
Institute of Science Education and Research Tirupati
<EMAIL_ADDRESS>
###### Abstract

Quantum Chromodynamics (QCD) predicts that at sufficiently high temperature
($T$) and/or baryon chemical potential ($\mu_{B}$), the state of matter is in
the form of quarks and gluons, which are no longer confined within hadrons.
This deconfined state of matter is known as the Quark-Gluon Plasma (QGP). The
goal of relativistic heavy-ion collision experiments is to create such a hot
and dense state of matter and study its properties. Measurements of identified
particle spectra in Au+Au collisions provide information on the bulk
properties, such as integrated yield ($dN/dy$), average transverse momenta
($\langle p_{T}\rangle$), particle ratios, and freeze-out parameters of the
medium produced. The systematic study of bulk properties sheds light on the
particle production mechanism in these collisions. Also, the centrality
dependence of the freeze-out parameters provides an opportunity to explore the
QCD phase diagram.

In this talk, we will present the transverse momentum spectra of identified
hadrons ($\pi^{\pm}$, $K^{\pm}$, $p$, and $\bar{p}$) at mid-rapidity
($|y|<0.1$) in Au+Au collisions at $\sqrt{s_{\mathrm{NN}}}=54.4$ GeV. The
centrality dependence of $dN/dy$, particle ratios, and kinetic freeze-out
parameters will also be presented, and their physics implications will be
discussed. In addition, we will compare our results with previously published
results at other collision energies.

##  1 Introduction

Quantum Chromodynamics (QCD) predicts the formation of the Quark-Gluon Plasma
(QGP), a new state of matter, in heavy-ion collisions at high energy density
or temperature [1]. Studying transverse momentum spectra in heavy-ion
collisions provides crucial information on QGP bulk properties, contributing
to our understanding of the QCD phase diagram, particle production mechanisms,
and freeze-out properties of the created medium. In this report, we present
the transverse momentum spectra of identified hadrons in Au+Au collisions at
$\sqrt{s_{NN}}=54.4$ GeV using the Time Projection Chamber (TPC) and Time of
Flight (TOF) detectors at STAR.

##  2 Results and Discussions

Figure 1: $\pi^{-}/\pi^{+}$, $K^{-}/K^{+}$, and $\bar{p}/p$ ratios at mid-
rapidity ($|y|<0.1$) in 0–5% Au+Au collisions at $\sqrt{s_{NN}}$ = 7.7–200
GeV. The uncertainties are statistical and systematic added in quadrature.

Figure 1 shows particle ratios ($\pi^{-}/\pi^{+}$, $K^{-}/K^{+}$, and
$\bar{p}/p$) in the most central (0-5%) collisions as a function of collision
energy. At lower beam energies, the $\pi^{-}/\pi^{+}$ ratios exceed unity due
to the contributions from resonance decays like $\Delta$ baryons. The
$K^{-}/K^{+}$ ratios show an increasing trend with increasing $\sqrt{s_{NN}}$
and approaches unity at higher beam energies, signifying the associated
production of $K^{+}$ at lower energies. The $\bar{p}/p$ ratios increase with
increasing $\sqrt{s_{NN}}$ but approach unity at the highest RHIC energy,
indicating stronger baryon stopping at lower energies. The 54.4 GeV results
follow the trend shown from previous measurements [2] of AGS, SPS, RHIC, and
LHC.

Figure 2: Variation of $T_{kin}$ with $\langle\beta\rangle$ for various
centralities in different collision energies.

A simultaneous fit to the $p_{T}$ spectra of $\pi$, $K$, $p$, and their
antiparticles was performed in different centrality intervals for Au+Au
collisions at $\sqrt{s_{NN}}=54.4$ GeV using the blast-wave model [3, 4] to
study the kinetic freeze-out properties of the medium. Figure 2 shows that as
we move from central to peripheral collisions, there is a decrease in
transverse flow velocity ($\langle\beta\rangle$) and an increase in kinetic
freeze-out temperature ($T_{kin}$), consistent with the expectation of a
shorter lived fireball towards peripheral collisions [5].

## References

  * [1] Krishna Rajgopal and Frank Wilczek  “The Condensed Matter Physics Of QCD”  In _At The Frontier of Particle Physics_ WORLD SCIENTIFIC, 2001, pp. 2061–2151  DOI: 10.1142/9789812810458˙0043
  * [2] L. Adamczyk et al.  “Bulk properties of the medium produced in relativistic heavy-ion collisions from the beam energy scan program”  In _Phys. Rev. C_ 96 American Physical Society, 2017, pp. 044904  DOI: 10.1103/PhysRevC.96.044904
  * [3] Ekkard Schnedermann, Josef Sollfrank and Ulrich Heinz  “Thermal phenomenology of hadrons from 200A GeV S+S collisions”  In _Phys. Rev. C_ 48 American Physical Society, 1993, pp. 2462–2475  DOI: 10.1103/PhysRevC.48.2462
  * [4] D. Teaney, J. Lauret and E.. Shuryak  “A Hydrodynamic Description of Heavy Ion Collisions at the SPS and RHIC”, 2001  arXiv:nucl-th/0110037 [nucl-th]
  * [5] Ulrich W. Heinz  “Concepts of Heavy-Ion Physics”, 2004  arXiv:hep-ph/0407360 [hep-ph]


# Bellman Optimal Step-size Straightening of Flow-Matching Models

Bao Nguyen
VinUniversity<EMAIL_ADDRESS>
&Binh Nguyen
National University of Singapore<EMAIL_ADDRESS>
Viet Anh Nguyen
Chinese University of Hong Kong<EMAIL_ADDRESS>
###### Abstract

Flow matching is a powerful framework for generating high-quality samples in
various applications, especially image synthesis. However, the intensive
computational demands of these models, especially during the fine-tuning
process and sampling processes, pose significant challenges for low-resource
scenarios. This paper introduces Bellman Optimal Step-size Straightening
(BOSS) technique for distilling flow-matching generative models: it aims
specifically for a few-step efficient image sampling while adhering to a
computational budget constraint. First, this technique involves a dynamic
programming algorithm that optimizes the step sizes of the pretrained network.
Then, it refines the velocity network to match the optimal step sizes, aiming
to straighten the generation paths. Extensive experimental evaluations across
image generation tasks demonstrate the efficacy of BOSS in terms of both
resource utilization and image quality. Our results reveal that BOSS achieves
substantial gains in efficiency while maintaining competitive sample quality,
effectively bridging the gap between low-resource constraints and the
demanding requirements of flow-matching generative models. Our paper also
fortifies the responsible development of artificial intelligence, offering a
more sustainable generative model that reduces computational costs and
environmental footprints. Our code can be found at
https://anonymous.4open.science/r/DRL-8E88.

##  1 Introduction

There are impressive advancements in deep generative models in recent years,
which constitute an appealing set of approaches capable of approximating data
distributions and generating high-quality samples, as showcased in influential
works such as Ramesh et al. (2022); Saharia et al. (2022); Rombach et al.
(2022). They are primarily driven by a category of time-dependent generative
models that utilize a predefined probability path, denoted as
$\\{\pi_{t}\\}_{t\in[0,1]}$. This probability path constitutes a process that
interpolates between the initial noise distribution $\pi_{0}$ and the target
data distribution $\pi_{1}$. The training for these models can be broadly
characterized as a regression task involving a neural network function
$v_{\theta}$ which maps any input data $x$ at time $t\in[0,1)$ to a vector-
valued velocity quantity $v_{\theta}(x,t)$ by minimizing the expected loss
between $v_{\theta}(x,t)$ and a target ideal velocity $v_{t}(x)$:

$\mathcal{L}(\theta)\coloneqq\mathbb{E}_{t\in[0,1],~{}X_{t}\sim\pi_{t}}[\ell(v_{\theta}(X_{t},t),v_{t}(X_{t}))].$

The velocity network $v_{\theta}$ plays a crucial role in the generation of
samples from the interpolation process through the relationship:

$X_{1}=X_{0}+\int_{0}^{1}v_{\theta}(X_{t},t)\mathrm{d}t,$

which is the solution of the ODE
$\mathrm{d}X_{t}=v_{\theta}(x_{t},t)\mathrm{d}t$ with the boundary condition
$X_{t=0}=X_{0}$. A noteworthy class of algorithms that fits within this
framework includes denoising diffusion models (Ho et al., 2020; Song et al.,
2020a; b) and the more recent flow-matching/rectified-flow models (Liu et al.,
2022b; Lipman et al., 2022; Albergo & Vanden-Eijnden, 2022; Neklyudov et al.,
2023). These models extend the principles employed in training diffusion
models to simulation-free continuous normalizing flows (CNF, Chen et al.
2018).

The rectified flow framework (Liu et al., 2022b) is particularly attractive
because it fixes the suboptimal alignment between noises and images of
diffusion models by introducing a straight trajectory formula connecting them.
This leads to faster training and inference time than diffusion models, at
least empirically. The rectified flow framework also introduces a technique
called _reflow_ , which gradually rectifies the probability paths. It
significantly reduces the number of function evaluations needed for sampling
and thus belongs to the family of distillation methods.

However, the standard reflow technique proposed in Liu et al. (2022b) requires
a significant amount of computational budget: on a small-dimensional dataset
such as CIFAR10 $(32\times 32$ pixel images), at least retraining the
pretrained model for another 300,000 iterations (Liu et al., 2022b) to reach
FID (Fréchet Inception Distance, Heusel et al. 2017) of 4.85 for 1-step
generation. It takes approximately 200 days of A100 GPU for distilling models
for 1-step sampling on higher-dimensional scale datasets to achieve
competitive FID (Liu et al., 2023). Motivated to fix this problem, in this
work, we aim to do distillation of Rectified Flow model while satisfying the
following objective.

Given a pretrained flow-matching velocity $v_{\theta}$ and a target of $K$
number of function evaluations (NFEs), how can we adapt $v_{\theta}$ with a
proper sampling schedule to generate high-fidelity images while using only a
modest amount of computational resource?

Contributions. We propose BOSS, the Bellman Optimal Step-size Straightening
method, to adapt pretrained flow-matching models. Our proposal includes two
phases. The first phase seeks the optimal $K$-element sequence $\Delta^{*}$
for the initial model $v_{\theta}$. The second phase utilizes $\Delta^{*}$ to
retrain $v_{\theta}$ such that the retrained model $v_{\theta^{*}}$ performs
better. With the proposed procedure, we succeeded in straightening the
velocity network with just about 10,000 of retraining iteration consistently
outperforms the standard reflow regarding image quality. Quantitatively, our
procedure consistently achieves lower FID in unconditional image generation
with four different datasets.

#### Related works.

Our approach is directly related to existing works on improving the sampling
efficiency of diffusion and flow-matching models with training-based
algorithms. In Salimans & Ho (2021), the authors proposed an approach to
enhance the sampling speed of unguided diffusion models through iterative
distillation. This is extended to the case of guided diffusion models in Meng
et al. (2023) that uses the distilled training of classifier-free guided
diffusion models. Recent work by ref:song23a proposes a framework that learns
a model capable of mapping any point at any time to the trajectory’s starting
point, called the _consistency model_. After submitting this work, we
discovered a concurrent study by Li et al. (2023). This work pointed out that
using uniform step size is suboptimal for diffusion model sampling and instead
used evolutionary algorithms to search for the optimal step size and score
network architectures, with the FID score being the optimized metric.

Within the context of Rectifed Flow/Flow Matching, Liu et al. (2022b) proposed
a reflow method that uses retraining to straighten the probability sampling
path. This results in a low NFE sampling with favorable image quality. The
recent work of Liu et al. (2023) takes this framework to a larger scale,
demonstrating impressive results on high-resolution image datasets. However,
both rely on computational intensive retraining procedures, which we improve
in our work.

The other direction that aims to accelerate the sampling process of
diffusion/flow matching models is training-free samplers (Song et al., 2020a;
Bao et al., 2022; Song et al., 2020b; Liu et al., 2022a; Tachibana et al.,
2021; Zhang & Chen, 2022; Karras et al., 2022; Lu et al., 2022; Zheng et al.,
2023). Although required no additional training step, these works mainly
relied on the properties of the SDE/probability flow ODE dynamics to propose
heuristic solvers/diffusion noise schedulers. Therefore, verifying whether the
proposed sampling step sizes are optimal is hard.

##  2 Background

Suppose we are given a (pretrained) model $v_{\theta}(X_{t},t)$ with parameter
$\theta$, which is an estimator of the function $v$ from an ordinary
differentiable equation (ODE) on the span $t\in[0,1]$:

$\mathrm{d}X_{t}=v(X_{t},t)\mathrm{d}t.$ (1)

Within the field of generative modeling with diffusion/flow matching models,
this dynamic system is called the _probability flow ODE_ (Song et al., 2020b;
Lipman et al., 2022). The estimator $v_{\theta}$ allows us to flow from the
distribution $\pi_{0}$ (noises) to the distribution $\pi_{1}$ (real images)
through the following equation:

$X_{1}=X_{0}+\int_{0}^{1}v_{\theta}(X_{t},t)\mathrm{d}t,$ (2)

where $X_{0}\sim\pi_{0}$ and $X_{1}\sim\pi_{1}$. In the context of our
problem, $X_{0}$ is observable. $X_{1}$ is only determined by the equation (2)
which is a deterministic process that for each $X_{0}=x_{0}$, there is only
value $X_{1}=x_{1}$ coupling with it through the following equation:

$x_{1}=x_{0}+\int_{0}^{1}v_{\theta}(x_{t},t)\mathrm{d}t.$

We are interested in the low-cost estimate of the integral
$\int_{0}^{1}v_{\theta}(x_{t},t)\mathrm{d}t$ with respect to $t$ over the
interval $[0,1]$. It is an essential concern when the cost of calculating the
velocity field is expensive. The amount of times calling $v_{\theta}$ is
denoted as the number of function evaluations (NFE).

First Order Sampling scheme. To solve for the integration that appears in the
sampling equation (2), it is necessary to invoke a numerical integrator that
uses discretized time steps. Any numerical integration scheme will induce
truncation errors, which can be quantified in two forms: first, when we have
the value at the previous time step $X_{\tau-\delta}$, the solver estimates
the subsequent true value $X_{\tau}$ as $\hat{X}_{\tau}$, causing a _local
truncation error_ $X_{\tau}-\hat{X}_{\tau}$. These local errors accumulate
over the number of intervals, eventually resulting in a cumulative error known
as the _global truncation error_. The most popular discretization scheme is
perhaps Euler’s method. Given a budget $K$ number of function evaluations
(NFEs), the Euler uniform sampling computes the interval $\Delta=1/K$, and the
sample successively

$x_{k/K}^{i}=x_{(k-1)/K}^{i}+v_{\theta}(x_{(k-1)/K}^{i},(k-1)\Delta)\times\Delta\quad\forall
k=1,\ldots,K,$ (3)

with the initial condition $x_{0}^{i}\sim\pi_{0}$. We denote this uniform
sampling scheme by $\mathcal{E}^{U}(K)$. Euler’s method with uniform step size
$\Delta$ has local truncation error $O(\Delta^{2})$, and global truncation
error $O(\Delta)$. One can generalize the Euler sampling with non-uniform
intervals by dividing the time domain $[0,1]$ into unequal intervals with
timestamps $0=\tau_{0}<\tau_{1}\ldots<\tau_{K}=1$, and sample successively

$x_{\tau_{k}}^{i}=x_{\tau_{k-1}}^{i}+v_{\theta}(x_{\tau_{k-1}}^{i},\tau_{k-1})\times(\tau_{k}-\tau_{k-1})\quad\forall
k=1,\ldots,K,$ (4)

with the initial condition $x_{0}^{i}\sim\pi_{0}$. The timestamps equivalently
determine the step-sizes $\tau_{k}-\tau_{k-1}$ for each sampling iteration. We
denote this scheme by $\mathcal{E}(\\{\tau_{0},\tau_{1},\ldots,\tau_{K}\\})$.
If the timestamps $\tau_{k}$ are equally spaced in $[0,1]$, then we obtain the
equivalence
$\mathcal{E}(\\{\tau_{0},\tau_{1},\ldots,\tau_{K}\\})\equiv\mathcal{E}^{U}(K)$.
Since the reflow procedure in Liu et al. (2022b) deals exclusively with
Euler’s method for being the fastest with a fixed computational budget, we
focus only on this method in our paper.

##  3 Optimal Sampling Step-sizes

The objective presented in Section 1 can be defined more rigorously as
follows. Given a pretrained model $v_{\theta}$ and an insignificant value of
$K$, we aim to find the optimal value $\theta^{*}$ and sequence $\Delta^{*}$
such that $\mathcal{E}(.,\Delta^{*})$ is a reasonable estimate for the
coupling sample $x_{1}$ of any $x_{0}$. This can be posed as an integer
optimization problem of finding the best schedule for sampling. Given a fixed
budget of $K$ NFEs, we find a schedule
$\\{\tau_{0},\tau_{1},\ldots,\tau_{K}\\}$ satisfying
$0=\tau_{0}<\tau_{1}<\ldots<\tau_{K}=1$ and that the associated Euler non-
uniform sampling scheme $\mathcal{E}(\\{\tau_{0},\tau_{1},\ldots,\tau_{K}\\})$
has minimal sampling error for the pretrained velocity $v_{\theta}$. In
Section 3.1, we describe our estimate of the sampling error for any valid
schedule. Section 3.2 presents an integer programming formulation to find the
optimal step sizes for sampling, and Section 3.3 provides a dynamic
programming algorithm to find the optimal schedule.

###  3.1 Sampling Error Estimation

Given any two arbitrary timestamps $0\leq t_{j}<t_{k}\leq 1$, we are
interested in estimating the local Euler truncation error, i.e., measuring the
discrepancy between the true value

$X_{t_{k}}=X_{t_{j}}+\int_{t_{j}}^{t_{k}}v_{\theta}(X_{t},t)\mathrm{d}t$

and the one-step Euler sampling value

$X_{t_{k}}^{\mathcal{E}}=X_{t_{j}}+v_{\theta}(X_{t_{j}},t_{j})\times(t_{k}-t_{j}),$

where $X_{t_{j}}$ is sampled from the distribution $\pi_{t_{j}}$ which is
induced by the initial distribution $\pi_{0}$ of $X_{0}$ and the ODE (1). This
local truncation error can be formalized as

$c_{jk}^{\mathrm{truncation}}\coloneqq\mathbb{E}_{X_{t_{j}}\sim\pi_{t_{j}}}\Big{[}\Big{\|}\int_{t_{j}}^{t_{k}}v_{\theta}(X_{t},t)\mathrm{d}t-v_{\theta}(X_{t_{j}},t_{j})(t_{k}-t_{j})\Big{\|}^{2}_{2}\Big{]}.$

Unfortunately, computing $c_{jk}^{\mathrm{true}}$ is computationally intensive
because of both the expectation operator and the integration. We instead
employ the following two simplifications:

  1. 1.

We fix the possible choice of time-stamps: for a sufficiently large number
$K^{\max}$, the anchoring timestamps are $\\{t_{k}\\}_{k=0,\ldots,K^{\max}}$
with $t_{k}=k/K^{\max}$. In doing so, we have restrained the space of all
possible sampling schedules to the combinations of finite anchoring timestamps
$\\{t_{k}\\}$. Later in the numerical experiments, we choose $K^{\max}=100$,
leading to the anchoring timestamps $\\{0,0.01,0.02,\ldots,0.99,1\\}$.

  2. 2.

We approximate the local truncation error $c_{jk}^{\mathrm{truncation}}$ for
any two anchoring timestamps $t_{j}<t_{k}$ by a sample average estimator
$c_{jk}$ that is constructed as follows:

$c_{jk}=\frac{1}{N}\sum_{i=1}^{N}c_{jk}^{i},\quad\mathrm{where}\quad
c_{jk}^{i}=\|x^{i}_{t_{k}}-x^{i}_{t_{j}}-v_{\theta}(x^{i}_{t_{j}},t_{j})\times(t_{k}-t_{j})\|^{2}_{2},$
(5)

where for sample $i$, the noise $x_{0}^{i}$ is drawn from $\pi_{0}$, and
$x^{i}_{t_{j}}$ and $x^{i}_{t_{k}}$ are extracted from Euler uniform sampling
path starting from $x_{0}^{i}$, taken at time $t_{j}$ and $t_{k}$,
respectively.

Figure 1: An example with $K^{\max}=5$ to illustrate the computation of the
sampling error.

Figure 1 illustrates how the sampling error $c_{jk}^{i}$ are calculated for
the simple case with $K^{\max}=5$ (or equivalently, with a uniform interval
$\Delta=0.2$). First, noise $x_{0}^{i}$ is drawn from $\pi_{0}$, and the blue
curve depicts the nonlinear trajectory following the ODE (1). The piecewise
linear trajectory is the path generated by the uniform Euler sampling with
$K^{\max}$ NFEs, leading to the observed trajectory
$\\{x_{k}^{i}\\}_{k=0,\ldots,K^{\max}}$. For a concrete example of computing
$c_{25}^{i}$, we measure the difference between the value of a one-step Euler
sampling from $t_{2}$ to $t_{5}$ with a stepsize $t_{5}-t_{2}=3\Delta$ to
obtain $x_{2}^{i}+v_{\theta}(x_{2}^{i},t_{2})\times 3\Delta$, and the observed
value $x_{5}^{i}$. Intuitively, we can view $c_{jk}$ as the difference between
the Euler one-step and the Euler $(k-j)$-step uniform sampling between $t_{j}$
and $t_{k}$.

One may observe that $c_{jk}$ is only an approximation of the true local
truncation error $c_{jk}^{\mathrm{truncation}}$ because $c_{jk}$ is computed
based on the Euler trajectory (red piecewise linear path in Figure 1), while
the truncation error $c_{jk}^{\mathrm{truncation}}$ should be computed based
on the nonlinear trajectory of the ODE (blue curve in Figure 1). One downside
of using $c_{jk}$ is that for any two consecutive timestamps $t_{j}$ and
$t_{j+1}$, we have $c_{j,(j+1)}=0$. This downside can be mitigated by taking
$K^{\max}$ sufficiently large. On the other hand, as $K^{\max}$ gets large,
calculating all the values $c_{jk}$ is computationally intensive because there
are, in total, $K^{\max}(K^{\max}-1)/2$ pair of timestamps whose errors are to
be computed. Nevertheless, we demonstrate empirically in Section 5 that even
when $c_{jk}$ is computed using a small number $N$ of samples, the resulting
optimal schedule already demonstrates a superior performance vis-à-vis
competing methods.

###  3.2 Integer Programming Formulation

As we now describe, finding the optimal sampling schedule can be formulated as
a network-flow-based problem (Ahuja et al., 1993). First, construct a graph of
$K^{\max}+1$ nodes; each node represents one timestamp, see Figure 2. There is
an edge connecting node $t_{j}$ to node $t_{k}$ if $t_{j}<t_{k}$, and this
edge is associated with a sampling error cost $c_{jk}$, computed in Section
3.1. For a target of $K$ NFEs, the optimal sampling schedule is a path that
traverses from the source node $t_{0}$ to the sink node $t_{K^{\max}}$ that is
comprised of exactly $K$ edges. This path can be recovered from the optimal
solution of the problem

$\begin{array}[]{cll}\min&\displaystyle\sum_{j=0}^{K^{\text{max}}-1}\sum_{k=j+1}^{K^{\text{max}}}c_{jk}z_{jk}\\\\[6.45831pt]
\operatorname{s.t.}&\sum_{j=0}^{K^{\text{max}}-1}\sum_{k=j+1}^{K^{\text{max}}}z_{jk}=K\\\\[6.45831pt]
&\sum_{k=1}^{K^{\text{max}}}z_{0k}=1,\quad\sum_{j=0}^{K^{\text{max}}-1}z_{jK^{\text{max}}}=1\\\\[6.45831pt]
&\sum_{k=0}^{j-1}z_{kj}=\sum_{k=j+1}^{K^{\text{max}}}z_{jk}&\forall
j\in\llbracket 1,K^{\text{max}-1}\rrbracket\\\ &z_{jk}\in\\{0,1\\}&\forall
0\leq j<k\leq K^{\max}.\end{array}$ (6)

Above, $z_{jk}$ is a binary decision variable, $z_{jk}=1$ if the path takes a
one-step sampling from time $t_{j}$ to time $t_{k}$. The objective function of
(6) is to minimize the path’s accumulated sampling error, which approximates
the global truncation error of the Euler sampling with the corresponding step
sizes. The first constraint indicates that the path should consist of exactly
$K$ edges; the second constraint imposes that $t_{0}$ and $t_{K^{\max}}$ are
the source and sink nodes, respectively. Finally, the last set of constraints
is the flow conservation on each intermediary node between $t_{0}$ and
$t_{K^{\max}}$.

Figure 2: A network flow formulation to find the optimal sampling schedule for
image generation. Time $t_{0}=0$ represents noise, while $t_{K^{\max}}=1$ is
the terminal data (images). Each discretized timestamp is represented by a
node, with edges reflecting the one-dimensional flow of time from noise to
image. The cost $c_{jk}$ associated with each edge is the sampling error
estimate, measured by the average difference between the Euler one-step and
the Euler $(k-j)$-step sampling between $t_{j}$ and $t_{k}$, see Section 3.1.

###  3.3 Dynamic Programming Algorithm

While the integer programming problem (6) can be solved using commercial
solvers such as GUROBI or using network flow algorithms (see Skiena (2008, §
6) for an example), there are practical cases in which we need to find optimal
paths for multiple values of the budget $K$ NFEs. A convenient way to address
this computation is to leverage a dynamic programming formulation, which
successively builds up the error-to-go function at each node and for each
number of remaining NFEs. To this end, for any timestamp $t_{\widehat{j}}$ and
any number of remaining NFEs $\widehat{k}\in\llbracket 1,K^{\max}\rrbracket$,
we define the error-to-go function as

$V({\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\widehat{k}}})\coloneqq\left\\{\begin{array}[]{cll}\min&\displaystyle\sum_{j={\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}}^{K^{\text{max}}-1}\sum_{k=j+1}^{K^{\text{max}}}c_{jk}z_{jk}\\\\[6.45831pt]
\operatorname{s.t.}&\sum_{j={\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}}^{K^{\text{max}}-1}\sum_{k=j+1}^{K^{\text{max}}}z_{jk}=\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\widehat{k}}\\\\[6.45831pt]
&\sum_{k={\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}+1}^{K^{\text{max}}}z_{{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}k}=1,\quad\sum_{j={\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}}^{K^{\text{max}}-1}z_{jK^{\text{max}}}=1\\\\[6.45831pt]
&\sum_{k={\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}}^{j-1}z_{kj}=\sum_{k=j+1}^{K^{\text{max}}}z_{jk}&\forall
j\in\llbracket{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}+1,K^{\text{max}-1}\rrbracket\\\
&z_{jk}\in\\{0,1\\}&\forall{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}\leq
j<k\leq K^{\max}.\end{array}\right.$ (7)

The error-to-go
$V({\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\widehat{k}}})$
is the minimal sampling error accumulated from time
$t_{{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}{\widehat{j}}}}$
to the terminal time $t_{K^{\max}}=1$ using exactly
$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\widehat{k}}$
NFEs. It is easy to see that the optimal value of problem (6) is equal to
$V(0,K)$. At initialization, we set the base case

$V(\widehat{j},1)=c_{\widehat{j}K^{\max}}\quad\forall~{}\widehat{j}\in\llbracket
1,K^{\max}\rrbracket.$

The dynamic programming update step rolls backward following

$\forall k=2,\ldots,K^{\max}:\qquad V(j,k)=\min_{j<\widehat{j}\leq
K^{\max}}c_{j\widehat{j}}+V(\widehat{j},k-1).$ (8)

The final output of the dynamic programming algorithm is the error function
$V$, and one can assess the Bellman optimal schedule by tracing the minimizing
path following (8) for each value $K$ of NFEs.

##  4 Straightening Flows with Bellman Step-size

Figure 3: Continued example following Figure 1 for straightening with $K=2$
NFEs, evaluated at time $t_{0}=0$ and time $t_{2}=0.4$. Blue arrows are
velocity vectors given by the pretrained model, purple arrows following the
dashed lines are the ideal straight path. The straightening procedure in
Section 4 aims to align the blue arrows towards purple arrows. Arrows
illustrate directions and are not drawn with proper scale.

Given the Bellman optimal step-sizes, we describe a piecewise linear
straightening of the sampling curve. The straightening procedure aims to re-
align the velocity network $v_{\theta}$ to reduce the accumulated sampling
error at the terminal timestamps. For a fixed number of NFEs $K$, let
$\\{\tau_{0},\ldots,\tau_{K}\\}$ be the optimal timestamps found in Section 3
with $\tau_{0}=0$ and $\tau_{K}=1$, which corresponds to $K$ stepsizes defined
by $\tau_{k}-\tau_{k-1}$ for $k=1,\ldots,K$. We now modify the network weights
to straighten the sampling path on each interval $[\tau_{k},\tau_{k+1}]$. An
intuitive explanation for the straightening procedure is illustrated in Figure
3: here, suppose that $K=2$, and the optimal schedule is
$\\{\tau_{0}=0,\tau_{1}=0.4,\tau_{2}=1\\}$. For the sample $i$ drawn in Figure
3, the Bellman sampling induces a two-piece linear path $x_{0}^{i}\to
x_{2}^{i}\to x_{5}^{i}$ (dashed line). If the velocity vectors evaluated at
$x_{0}^{i}$ and $x_{2}^{i}$ align with the dashed line, then the Bellman
optimal Euler sampling with $K=2$ incurs zero loss compared to the Euler
uniform sampling with $K^{\max}=5$. This motivates the following alignment
procedure:

$\min_{\theta}~{}\mathbb{E}_{X_{0}\sim\pi_{0}}\Big{[}\sum_{k=0}^{K-1}\Big{\|}v_{\theta}(X_{\tau_{k}}^{\mathcal{E}(K^{\max})},\tau_{k})-\frac{X_{\tau_{k+1}}^{\mathcal{E}(K^{\max})}-X^{\mathcal{E}(K^{\max})}_{\tau_{k}}}{\tau_{k+1}-\tau_{k}}\Big{\|}_{2}^{2}\Big{]},$

where $X_{\tau_{k}}^{\mathcal{E}(K^{\max})}$ is obtained by the Euler uniform
sampling with $K^{\max}$ NFEs of the initial condition $X_{0}\sim\pi_{0}$.
Replacing the expectation with $n$ empirical paths obtained by
$\mathcal{E}(K^{\max})$, we have the sample averaging optimization problem for
straightening

$\min_{\theta}~{}\frac{1}{n}\sum_{i=1}^{n}\sum_{k=0}^{K-1}\|v_{\theta}(x^{i}_{\tau_{k}},\tau_{k})-\frac{x^{i}_{\tau_{k+1}}-x^{i}_{\tau_{k}}}{\tau_{k+1}-\tau_{k}}\|_{2}^{2}.$
(9)

We straighten the velocity model using a stochastic gradient descent algorithm
to solve the above problem.

##  5 Numerical Experiments

#### Settings.

We evaluate our methods on unconditioned image generation tasks. In
particular, we use the CIFAR10 (Krizhevsky et al., 2009) and three high-
resolution (256x256) datasets CelebA-HQ (Karras et al., 2018), LSUN-Church,
LSUN-Bedroom (Yu et al., 2015), and AFHQ-CAT. We take the checkpoints of
pretrained velocity networks $v$ from the official
implementation111https://github.com/gnobitab/RectifiedFlow/ of Rectified Flow
(Liu et al., 2022b), which is based on the U-Net architecture of DDPM++ (Song
et al., 2020b). If not mentioned otherwise, we evaluate the sample schemes
with NFE=$\\{4,6,8\\}$. The quality of image samples is with Frechet inception
distance (FID) score (Heusel et al., 2017).

#### Baselines.

A comparison between Bellman optimal step sizes and the conventional
first/second order methods using uniform step sizes is presented in Section
5.1. We also include an adaptive strategy, the Runge-Kutta method of order
5(4) from SciPy (Virtanen et al., 2020). In Section 5.2, our finetuning
procedure, Bellman Optimal Step-size Straightening (BOSS) is compared with two
baselines including the Uniform-Reflow, and Distill-k-Reflow introduced in
(Liu et al., 2022b).

###  5.1 Improvements of first order and second order sampling scheme using
Bellman Optimal Stepsize

First, we benchmark sampling on pretrained models with Euler’s method with
uniform stepsize and with Bellman optimal stepsize, calculated following
Section 3.3. The quantitative results are demonstrated in Table 1. The FID is
clearly much lower for samples generated by Euler’s method with the Bellman
step, which shows that with Bellman’s optimal step size, the generated images
are of much higher quality in general. Specifically, for the three larger
dimension datasets (256x256), sampling with Bellman steps can help drastically
reduce FID compared to the uniform step size. This trend is also reflected in
Figure 5 that shows our qualitative results.

Table 1: FID ($\downarrow$) of Euler’s sampling method with uniform step size vs. Bellman optimal step size on unconditional image generation task across different datasets. Dataset | 4 NFE | 6 NFE | 8 NFE
---|---|---|---
Euler | Bellman | Euler | Bellman | Euler | Bellman
CIFAR10 | 51.95 | $\mathbf{47.57}$ | 25.69 | $\mathbf{23.35}$ | 16.82 | $\mathbf{15.74}$
CelebA-HQ | 158.95 | $\mathbf{92.03}$ | 127.01 | $\mathbf{72.54}$ | 109.42 | $\mathbf{49.80}$
LSUN-CHURCH | 106.94 | $\mathbf{80.91}$ | 53.85 | $\mathbf{45.09}$ | 34.74 | $\mathbf{33.22}$
AFHQ-CAT | 68.95 | $\mathbf{45.54}$ | 61.50 | $\mathbf{36.15}$ | 56.96 | $\mathbf{33.94}$
LSUN-BEDROOM | $84.35$ | $\mathbf{61.60}$ | $39.19$ | $\mathbf{35.35}$ | $32.15$ | $\mathbf{25.80}$

We also empirically analyze the effect of Bellman optimal stepsizes on popular
ODE solvers including Euler and Heun (second-order version). To avoid
confusing between schemes, we denote compared methods as follows:

  * •

Uniform Euler and Uniform Heun are the conventional Euler and Heun methods
that use uniform sampling steps.

  * •

Bellman Euler and Bellman Heun are two variants of the above methods, but
using our proposed Bellman step sizes.

  * •

RK45 is an adaptive strategy, the Runge-Kutta method of order 5(4) from Scipy.

The quantitative results are displayed in Figure 4. It is evident that
employing Bellman optimal stepsizes can significantly improve FID scores
(image quality) compared to using a uniform stepsize across all pretrained
models on four distinct datasets. The Bellman Heun methods approach the
performance of RK45 with substantially fewer NFEs. To elaborate, the Bellman
Heun method achieves approximately a 1% gap compared to RK45 with just 20
sampling steps and fully recovers the performance of RK45 with 50 steps.

(a)

(b)

(c)

(d)

(e)

Figure 4: The FID score of sampling methods with different numbers of function
evaluations (step sizes). Images generated by samplers using Bellman step
sizes clearly show lower FID than conventional ones that use uniform step
sizes. Note that Uniform Heun and Bellman Heun are second-order sampling
methods that use twice the NFEs.

###  5.2 Effects of Reflow with Bellman Step-size

After calculating the Bellman optimal step size, we follow the procedure
described in Section 4 to straighten the probability path. The results, seen
in Table 2, suggest that BOSS performs almost equally with the reflow
procedure on CIFAR10 but markedly better on the other four higher-dimension
datasets. This is consistent with the visible improvements in sampled image
quality observed in Figure 5.

Table 2: FID ($\downarrow$) of different retraining methods on the unconditional image generation task across different datasets. Dataset | Distill-6-Reflow | BOSS-6 | Uniform-Reflow
---|---|---|---
CIFAR10 | $4.35$ | $4.80$ | $\mathbf{4.33}$
CelebA-HQ | $34.56$ | $\mathbf{18.67}$ | $43.57$
LSUN-CHURCH | $34.52$ | $\mathbf{17.43}$ | $40.45$
AFHQ-CAT | $46.24$ | $\mathbf{26.10}$ | $51.24$
LSUN-BEDROOM | $41.17$ | $\mathbf{18.45}$ | $45.13$
|  |  |
---|---|---|---
|  |  |
|  |  |
|  |  |
(a) Euler (6 NFEs) | (b) Bellman (6 NFEs) | (c) BOSS (6 NFEs) | (d) RK45 (208 NFEs)
Figure 5: Qualitative results on unconditional image generation task. From
first to last row: CelebA-HQ-256/LSUN Bedroom/LSUN Church/AFHQ dataset.
(a)-(b): Comparisons of Euler step sizes between uniform (a) and the Bellman
optimal step sizes (b); (c)-(d): Comparisons of BOSS retraining and Runge-
Kutta-45 sampling. Notice our proposed BOSS sampling has comparably similar
visual quality to RK45 while requiring only 6 NFEs, compared to 208 NFEs of
RK45.

##  6 Conclusions

This paper proposed BOSS, the Bellman Optimal Step-size Straightening method,
to adapt pretrained flow-matching models under low computational resource
constraints. Our method consists of two phases: first, find an optimal
sampling step sizes for the pretrained model, then straighten out the velocity
network on each interval of the sampling schedule. We demonstrate empirically
that BOSS performs competitively in adapting pretrained models in the image
generation task. Similar to training-based samplers for diffusion and flow
matching models, a limitation of our method is the additional training cost to
output the optimal sample step sizes. There are many potential future
extensions to our proposed framework to distill a guided velocity network,
similarly to Meng et al. (2023), or a computationally-cheaper algorithm for
calculating the Bellman sampling step sizes.

## References

  * Ahuja et al. (1993) R.K. Ahuja, T.L. Magnanti, and J.B. Orlin.  _Network Flows: Theory, Algorithms, and Applications_.  Prentice Hall, 1993.
  * Albergo & Vanden-Eijnden (2022) Michael Samuel Albergo and Eric Vanden-Eijnden.  Building normalizing flows with stochastic interpolants.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Bao et al. (2022) Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang.  Analytic-DPM: an analytic estimate of the optimal reverse variance in diffusion probabilistic models.  _arXiv preprint arXiv:2201.06503_ , 2022.
  * Chen et al. (2018) Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Heusel et al. (2017) Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  GANs trained by a two time-scale update rule converge to a local Nash equilibrium.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Karras et al. (2018) Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen.  Progressive growing of GANs for improved quality, stability, and variation.  In _International Conference on Learning Representations_ , 2018.
  * Karras et al. (2022) Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the design space of diffusion-based generative models.  _Advances in Neural Information Processing Systems_ , 35:26565–26577, 2022.
  * Krizhevsky et al. (2009) Alex Krizhevsky et al.  Learning multiple layers of features from tiny images.  _cs.toronto.edu_ , 2009.
  * Li et al. (2023) Lijiang Li, Huixia Li, Xiawu Zheng, Jie Wu, Xuefeng Xiao, Rui Wang, Min Zheng, Xin Pan, Fei Chao, and Rongrong Ji.  Autodiffusion: Training-free optimization of time steps and architectures for automated diffusion model acceleration.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 7105–7114, 2023.
  * Lipman et al. (2022) Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le.  Flow matching for generative modeling.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Liu et al. (2022a) Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao.  Pseudo numerical methods for diffusion models on manifolds.  _arXiv preprint arXiv:2202.09778_ , 2022a.
  * Liu et al. (2022b) Xingchao Liu, Chengyue Gong, and Qiang Liu.  Flow straight and fast: Learning to generate and transfer data with rectified flow.  _arXiv preprint arXiv:2209.03003_ , 2022b.
  * Liu et al. (2023) Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu.  Instaflow: One step is enough for high-quality diffusion-based text-to-image generation.  _arXiv preprint arXiv:2309.06380_ , 2023.
  * Lu et al. (2022) Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu.  DPM-solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps.  _Advances in Neural Information Processing Systems_ , 35:5775–5787, 2022.
  * Meng et al. (2023) Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans.  On distillation of guided diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 14297–14306, 2023.
  * Neklyudov et al. (2023) Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani.  Action matching: Learning stochastic dynamics from samples.  In _Proceedings of the 40th International Conference on Machine Learning_ , 2023.
  * Parmar et al. (2022) Gaurav Parmar, Richard Zhang, and Jun-Yan Zhu.  On aliased resizing and surprising subtleties in GAN evaluation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022.
  * Ramesh et al. (2022) Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen.  Hierarchical text-conditional image generation with CLIP latents.  _arXiv preprint arxiv:2204.06125_ , 7, 2022.
  * Rombach et al. (2022) Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  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.
  * Saharia et al. (2022) Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al.  Photorealistic text-to-image diffusion models with deep language understanding.  _Advances in Neural Information Processing Systems_ , 35:36479–36494, 2022.
  * Salimans & Ho (2021) Tim Salimans and Jonathan Ho.  Progressive distillation for fast sampling of diffusion models.  In _International Conference on Learning Representations_ , 2021.
  * Skiena (2008) Steven S. Skiena.  _The Algorithm Design Manual_.  Springer, 2008.
  * Song et al. (2020a) Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  In _International Conference on Learning Representations_ , 2020a.
  * Song et al. (2020b) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  In _International Conference on Learning Representations_ , 2020b.
  * Tachibana et al. (2021) Hideyuki Tachibana, Mocho Go, Muneyoshi Inahara, Yotaro Katayama, and Yotaro Watanabe.  Quasi-Taylor samplers for diffusion generative models based on ideal derivatives.  _arXiv preprint arXiv:2112.13339_ , 2021.
  * Virtanen et al. (2020) Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, et al.  SciPy 1.0: Fundamental algorithms for scientific computing in Python.  _Nature Methods_ , 17(3):261–272, 2020.
  * Yu et al. (2015) Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao.  Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop.  _arXiv preprint arXiv:1506.03365_ , 2015.
  * Zhang & Chen (2022) Qinsheng Zhang and Yongxin Chen.  Fast sampling of diffusion models with exponential integrator.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Zheng et al. (2023) Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu.  DPM-solver-v3: Improved diffusion ODE solver with empirical model statistics.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.

##  Appendix A Details of Experiments

We use the following checkpoints that are downloaded from the GitHub folder
222https://github.com/gnobitab/RectifiedFlow/:

  * •

CIFAR10: at iteration 800,000;

  * •

CelebA-HQ-256: at iteration 1,000,000;

  * •

LSUN-Church: at iteration 1,200,000.

  * •

LSUN-Bedroom: at iteration 1,000,000.

  * •

AFHQ-Cat: at iteration 1,000,000.

The pre-trained models are finetuned by these methods in 12000 iterations. An
iteration is the passing and backpropagation process for a batch including 15
samples. Due to the similar cost of training between finetuning methods, we
report the average GPU hours consumed on each pre-trained model up to 12000
iterations, using NVIDIA RTX A5000.

  * •

CIFAR10: 3.56 training hours.

  * •

CelebA-HQ: 10.35 training hours.

  * •

LSUN-CHURCH: 13.43 training hours.

  * •

LSUN-BEDROOM: 14.23 training hours.

  * •

AFHQ-CAT: 9.30 training hours.

With this limited budget of resources, the proposed method, BOSS, achieves
significantly better performance than other methods in terms of FID score. The
value of $N$ in Equation (5) and $K^{\max}$ are fixed at $100$, and $100$ in
all experiments if not mentioned. This setup $N=100$ means we only use one
batch sampling to calculate the truncation errors between timestamps in
equation (5). It is worth noting that this setup highlights the low-cost and
limited-resources requirement of our proposal.

#### FID Calculation.

All the FID metrics of related works are either cited from previous baselines
or are calculated (when there are no such figures reported) based on the
Clean-FID paper (Parmar et al., 2022), which unifies the FID calculation to
make a fair comparison between papers. These four datasets were downloaded
following the instructions from their original papers. We then create the
stats file by the clean-fid project. The FID score is then calculated based on
50000 generated images and the stats file.

##  Appendix B Description of the Dynamic Programming Algorithm

This section presents the pseudocode in Algorithm 1 for the practical
implementation of the dynamic programming algorithm designed to determine the
Bellman optimal stepsizes. The algorithm takes a cost matrix, denoted as $c$,
as input, where $c_{jk}$ is computed using Equation (5) and the specified
Number of Function Evaluations (NFEs). The most resource-intensive aspect of
this code is the nested loop responsible for calculating $\kappa(j,k)$,
incurring a time complexity of $O((K^{\max})^{2}\times K)$. The choice of
$K^{\max}$ is crucial, aiming for the Euler sampling method with $K^{\max}$
stepsizes to precisely replicate the trajectory of the Ordinary Differential
Equation (ODE). Typically, $K^{\max}$ falls within the range of 100 to 1000,
ensuring accuracy. Given this range for $K^{\max}$ and the variable $K$
ranging from 2 to 1000, the algorithm executes within milliseconds in all
scenarios.

Algorithm 1 Minimum Cost Path Computation

Input: Cost matrix $c(j,k)=c_{jk}(j,k=0,\ldots,K^{\max})$ , the target NFEs
$K$.

Set $\kappa(i,j)\leftarrow+\infty\quad\forall 0\leq i\leq K^{\max},0\leq j\leq
K$

Set $\kappa(i,1)\leftarrow c(i,K^{\max})\quad\forall 0\leq i\leq K^{\max}$

for $k=2$ to $K$ do

for $j=0$ to $K^{\max}-1$ do

for $i=j+1$ to $K^{\max}-1$ do

$\kappa(j,k)\leftarrow\min\\{\kappa(j,k),c(j,i)+\kappa(i,k-1)\\}$

end for

end for

end for

Initialize $\psi\leftarrow[0]$, $\omega\leftarrow 0$

for $k=K$ to $1$ do

for $j=\omega+1$ to $K^{\max}$ do

if $\kappa(\omega,k)=c(\omega,j)+\kappa(j,k-1)$ then

Append $j$ to $\psi$ and set $\omega\leftarrow j$

break

end if

end for

end for

Append $K^{\max}$ to $\psi$.

return $\psi,\kappa(0,K)$

##  Appendix C Empirical analysis about Bellman Optimal Step-sizes

###  C.1 A common trend in bellman optimal stepsizes for pretrained models on
different datasets

We plot the Bellman Optimal Step sizes in Figure 6. It shows a common trend of
sampling with small initial steps and then larger step sizes for intermediate
iterations. At $K=6$, we can observe that the last step is smaller than the
penultimate step, hinting that the sampling process aims to take smaller final
steps to refine the output. This refining trend is more evident for $K=8$.

Figure 6: Optimal Bellman step sizes for CIFAR10 and CelebA at $K=2,4$ and
$6$. One can identify a common pattern of smaller steps at the beginning and
the end of the sampling procedure

###  C.2 Empirical Evidence for the Step Size Trend

This section aims to experimentally explain the trend of smaller stepsizes at
the beginning and the end of the sampling procedure. Given a velocity network
$v_{\theta}$, we empirically estimate the curvature using the following
procedure:

  1. 1.

Sample $N$ noise instances $x_{0}^{i}$ i.i.d. from a Gaussian distribution (we
set $N\approx 1000$).

  2. 2.

Forward each noise instance using $v_{\theta}$ to obtain $x_{t}^{i}$ for
$t=1,\ldots,K^{\max}$.

  3. 3.

Compute the local curvature for each trajectory:

$\mathrm{Curv}_{t}^{i}=\|x_{t}^{i}-\frac{x_{t+1}^{i}+x_{t-1}^{i}}{2}\|_{2}^{2}\quad\text{for
}t=2,\ldots,K^{\max}-1.$

  4. 4.

Calculate the average curvature:

$\widehat{\mathrm{Curv}}_{t}=\frac{1}{N}\sum_{i=1}^{N}\mathrm{Curv}_{t}^{i}.$

Subsequently, we plot $\widehat{\mathrm{Curv}}_{t}$ and investigate whether
the Bellman step size coincides with the straightness of the curve, as
demonstrated in Figure 7. The curvature of all pretrained rectified models is
significant at timestamps close to zero (near the space of noises) and one
(near the space of real images). This observation indicates a variant stepsize
trajectory, notably at initial and final timestamps. It can be intuitively
explained that additional steps are required to match the high-curvature
region accurately. Consequently, the outcome of our proposal is sensible, as
it correctly identifies the high curvature levels at the beginning and end of
the sampling trajectories.

(a)

(b)

(c)

(d)

(e)

Figure 7: Curvature measurement (blue curve) and Bellman optimal timestamps
(black horizontal line) for different datasets. Note that the timestamps are
denser at regions with higher estimated curvature

###  C.3 Finding optimal step-sizes is a light-weight process

In this section, we elaborate further on the efficiency of our framework. This
efficiency arises because we only need to pass a single batch of noise through
the forward process to obtain the values for each intermediate timestamp.
Subsequently, we calculate the local truncation error for any two timestamps.
These local truncation errors have in total $K^{\max}\times(K^{\max}-1)/2$,
and they can be efficiently stored without requiring large memory. The dynamic
programming involved in this process is also time-efficient, as elaborated in
the Appendix B. For added credibility, we provide a running time of the entire
stepsizes calculation process with NFE = 10, detailing the running time of
each component across all our datasets in Table 3. All running times are for
the Nvidia A5000 which is an old-generation GPU launched in April 2021. The
whole process takes around 115 seconds to complete for the 256x256 datasets.

Table 3: Running time (seconds) for different datasets | CIFAR10 | CelebA-HQ | LSUN-CHURCH | LSUN-BEDROOM | AFHQ-CAT
---|---|---|---|---|---
One batch forward | 8.5977 | 47.7143 | 47.4024 | 48.4324 | 45.5489
Local truncation errors | 10.066 | 67.9143 | 67.5273 | 65.9875 | 66.7532
Dynamic programming | 0.0134 | 0.0138 | 0.0138 | 0.0139 | 0.0138
The whole process | 18.6771 | 115.6424 | 114.9435 | 114.4338 | 112.3159
Table 4: Bellman Optimal Step sizes for $K=2$ and $K=4$. The total sum of step sizes equals one. | $K=2$ | $K=4$
---|---|---
CIFAR10 | $[0.28,0.72]$ | $[0.12,0.22,0.31,0.35]$
CELEBA-HQ | $[0.13,0.87]$ | $[0.05,0.17,0.32,0.46]$
CHURCH | $[0.20,0.80]$ | $[0.08,0.21,0.34,0.37]$
Table 5: Bellman Optimal Step sizes for $K=6$ and $K=8$. The total sum of step sizes equals one. | $K=6$ | $K=8$
---|---|---
CIFAR10 | $[0.09,0.13,0.17,0.21,0.22,0.18]$ | $[0.07,0.09,0.12,0.14,0.17,0.17,0.14,0.10]$
CELEBA-HQ | $[0.03,0.08,0.17,0.26,0.28,0.18]$ | $[0.02,0.05,0.11,0.17,0.21,0.20,0.16,0.08]$
CHURCH | $[0.06,0.15,0.24,0.26,0.20,0.09]$ | $[0.04,0.09,0.15,0.17,0.18,0.17,0.13,0.07]$

The Bellman steps being far from uniform also means that the probability path
of the pretrained models is far from straight, and performing the
straightening operation would be beneficial.

We report in this section the Bellman optimal step-sizes obtained in Section
3.3. We want to focus on the case $K=8$ to see the common trend of sampling
step sizes. Initially, the sampling process takes small step sizes, possibly
for structural determination of the images. The step sizes become larger for
the intermediate steps. The last two step sizes show a decreasing trend: the
sampling process takes small step sizes at the end to refine and potentially
make the final output less noisy.

##  Appendix D The transfer of optimal stepsizes from a dataset to others

To verify the generalization of optimized stepsizes, we transferred the
optimized step sizes from LSUN-CHURCH to the pre-trained models on CelebA-HQ
and LSUN-BEDROOM. The FID scores obtained with 4, 6, and 8 NFEs for CelebA-HQ
resulting from this transfer are presented in Table 6.

Table 6: FID scores for CelebA-HQ with different methods and NFEs Method | 4 NFE | 6 NFE | 8 NFE
---|---|---|---
Uniform Euler | 158.95 | 127.01 | 109.42
Bellman Euler | 92.03 | 72.54 | 49.80
Bellman-transfer | 132.04 | 100.68 | 72.88

Uniform Euler uses uniform step sizes, Bellman Euler uses optimal step sizes
for CelebA-HQ, while Bellman-transfer uses the step sizes taken from LSUN-
CHURCH. Bellman Euler is still the optimal method. However, what is important
here is that Bellman-transfer is better than Uniform Euler. This hints that
there is a certain degree of transferability of the step sizes.333This is an
empirical claim; we do not impose any theoretical claim. In the empirical
realm, we can attribute this transferability to a comparable curvature pattern
exhibited by pretrained rectified models, as discussed in Section C.2.

The Table 7 is for the LSUN-BEDROOM dataset. We observe the same trend here,
confirming empirically that the step size has a certain degree of
transferability. Nevertheless, optimizing the step size using Bellman Euler
would still obtain the best performance.

Table 7: FID scores for LSUN-BEDROOM with different methods and NFEs Method | 4 NFE | 6 NFE | 8 NFE
---|---|---|---
Uniform Euler | 84.35 | 39.19 | 32.15
Bellman Euler | 61.60 | 35.35 | 25.80
Bellman-transfer | 70.23 | 38.01 | 29.14

##  Appendix E Additional Qualitative Results

Figure 8: Comparison between images generated from an identical noise with
different sampling methods and the number of stepsizes on the LSUN-CHURCH
dataset. Figure 9: Comparison between images generated from an identical noise
with different sampling methods and the number of stepsizes on the LSUN-
BEDROOM dataset.

(a)

(b)

Figure 10: Comparative qualitative outcomes of BOSS with NFE = 6. The image on
the right showcases the generated images referenced by RK45. Figure 11:
Uncurated images generated from the model fine-tuned by BOSS (NFE = 10,
FID=13.89)

(a)

(b)

(c)

(d)

(e)

(f)

Figure 12: Comparative qualitative outcomes of BOSS with different NFEs on the
LSUN-BEDROOM dataset.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 13: Samples from LSUN-BEDROOM 256x256. All corresponding samples use
the same initial noise.