AdaFlow: Imitation Learning with Variance-Adaptive Flow-Based Policies
Authors: Xixi Hu
qiang liu
Xingchao Liu
Bo Liu
Venue: NeurIPS 2024 poster (2024)
Keywords: Imitation Learning; Generative Model-Based Policy
Primary area: diffusion_based_models

Abstract: Diffusion-based imitation learning improves Behavioral Cloning (BC) on multi-modal decision-making, but comes at the cost of significantly slower inference due to the recursion in the diffusion process. It urges us to design efficient policy generators while keeping the ability to generate diverse actions. To address this challenge, we propose AdaFlow, an imitation learning framework based on flow-based generative modeling. AdaFlow represents the policy with state-conditioned ordinary differential equations (ODEs), which are known as probability flows. We reveal an intriguing connection between the conditional variance of their training loss and the discretization error of the ODEs.
With this insight, we propose a variance-adaptive ODE solver that can adjust its step size in the inference stage, making
AdaFlow an adaptive decision-maker, offering rapid inference without sacrificing diversity. Interestingly, it automatically reduces to a one-step generator when the action distribution is uni-modal. Our comprehensive empirical evaluation shows that AdaFlow achieves high performance with fast inference speed.

AdaFlow: Imitation Learning with Variance-Adaptive Flow-Based Policies
Xixi Hu, Bo Liu, Xingchao Liu, Qiang Liu
The University of Texas at Austin
(hxixi,bluu, xcliu, aqgangg utexas edu
Abstract
Diffusion-based imitation learning improves Behavioral Cloning (BC) on multi- modal decision-making, but comes at the cost of significantly slower inference due the recursion in the diffusion process. urges to design efficient policy generators while keeping the ability to generate diverse actions. To address this challenge, we propose AdaFlow, an imitation learning framework based on flow- based generative modeling. AdaFlow represents the policy with state-conditioned ordinary differential equations (ODEs), which are known as probability flows. We reveal an intriguing connection between the conditional variance of their training loss and the discretization error of the ODEs. With this insight, we propose a variance-adaptive ODE solver that can adjust its step size in the inference stage, making AdaFlow an adaptive decision-maker, offering rapid inference without sacrificing diversity. Interestingly, it automatically reduces to one-step generator when the action distribution is uni-mat mni-m Our comprehensive empirical evaluation shows that AdaFlow achieves high performance with fast inference speed.
Now-Variance State
High- Variunce State
Low
High
Diffusion Policy
Noise
Step 50
Step 100
Noise
Step 50
Step 100
AdaFlow
1 Step
State Variance
Action
Noise
Step
Step
Figure 1: AdaFlow is a fast imitation learning policy. It adaptively adjust the number of simulation steps when generating actions. For low variance states, it functions as one-step action generator. For high-variane states, employs more steps average generation speed close step per task completion.
accurate action generation. This adaptive approach enables AdaFlow to achieve
Introduction
Imitation Learning (IL) is widely adopted method in robot learning 21 In IL, an agent is given demonstration dataset from human expert finishing certain task, and the goal is for complete the task by learning from this dataset. IL is notably effective for learning complex, non-declarative motions, yielding remarkable successes training real robots
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Data
Prediction
X -2.5
x= 25
X
10
50-
y
X
X
Ground Truth
Predicted Variance
Diffusion Policy
20
AdaFlow (Euler)
AdaFlow
25
Illustrating the computation adaptivity of AdaFlow (orange) simple regression task. In the upper Diffusion Policy (DDIM) and AdaFlow predict y given æ, with deterministic portion of the image, bimodal y = iir when x > Both DDIM and AdaFlow the demonstration data y = when &
Figure 2:
well. However, the simulated ODE trajectory learned by Diffusion-Policc Policy with DDIM (red) is not straight no matter what x is. By contrast, the simulated ODE trajectory learned by AdaFlow with fixed step (bluc) is a straight line when the prediction is deterministic (z S 0), which means the generation be exactly done by one-step Euler discretization. simulation steps based on æ value according to estimated variance at x.
bottom, we show that AdaFlow can adaptively adjust the number of
The primary approach for IL is Behavioral Cloning (BC) 7-10,, where the agent is trained with supervised learning to acquire a deterministic mapping from states to actions. Despite its simplicity,
vanilla BC struggles to learn diverse behaviors in states with many possible actions [10 [1] 12).
To improve it, various frameworks have been proposed. For instance, Implicit BC [1] learns an
energy-based model for each state and searches the actions that minimize the energy with optimization
algorithms. Diffuser [13] [4] and Diffusion Policy [6] adopts diffusion models [T5] 16 to generate
diverse actions, which has become the default method for training on large-scale roboties data [7-
The computational cost of the learned policies at the execution stage is important for an framework in real-world deployment. Unfortunately, none of the previous frameworks enjoys both efficient inference and diversity. Although energy- models and diffusion models can generate multi- modal action distributions, they require recursive processes to generate the actions. These recursive processes usually involve tens (or even hundreds) of queries before reaching their stopping criteria.
In this paper, propose new IL framework, named AdaFlow, that learns dynamic generative policy that can autonomously adjust computation on the fly, thus cheaply outputs multi-modal action distributions to complete the task. AdaFlow is inspired by recent advancements in flow-based generative modeling (21/24]. We learn probability flows, which are essentially ordinary differential equations (ODEs), to represent the policies. These flows powerful generative models that precisely capture the complicated distributions, but similar to energy-basee models and diffusion models, they still require multiple recursive iterations simulate the ODEs in the inference stage.
AdaFlow differs from existing flow generative models like Rectified Flow 25 and Consistency Models (26), by utilizing the initially learned ODE to maintain low training and inference costs, and function as a one-step generator for deterministic target distributions. contrast, both of these methods require an additional distillation reflow 25 process to achieve fast inference. To improve the efficiency, we propose an adaptive ODE solver based on the finding that the simulation error of the ODE is closely related to the variance of the training loss at different states. We let the action generation model to output an additional variance scalar alongside the action it produces. During the execution of the policy, we change the step size according to the variance predicted at the current state. Equipping the flow-based policy with the proposed adaptive ODE solver, AdaFlow wisely allocates computational resources, yielding high efficiency without sacrificing the diversity provided by the flow-based generative models. Specifically, in states with deterministic action distributions, AdaFlow generates the action in one step as efficient naive BC.
Empirical results across multiple benchmarks demonstrate that AdaFlow achieve consistently good performance across success rate with high exccution efficiency. Specifically, contributions are:

---

We proposed AdaFlow generative model-hased policy for decision making tasks, capable of generating actions almost as quickly as single model inference pass.
We conducted comprehensive experiments across decision making tasks, including navigation and robot manipulation, utilizing benchmarks such as LIBERO [27] and RoboMimic [1]]. AdaFlow consistently outperforms existing state-of-the-ar models, de- spite requiring 10x less inference times to generate actions.
We offer a theoretical analysis of the overall error in action generation by AdaFlow, providing bound that ensures precision and reliability.
2 Related Work
Diffusion/Flow- Generative Models and Adaptive Inference Diffusion models 128 16
15 29 succeed in various applications, c.g., image/video generation 33) audio generation 34]
point cloud generation B3) etc.. However, numerical simulation of the diffusion processes typically involve hundreds of steps, resulting in high inference cost. Post-hoc samplers have been proposed to solve this issue 39 4 by transforming the diffusion process into
probability flow ODEs, yet they still use the same number of inference steps for different states. Although adaptive ODE solvers, such as adaptive step-size Lunge-KKutte [45), exist, they cannot significantly reduce the number of inference steps. In comparison, the adaptive sampling strategy of AdaFlow is specifically designed based on intrinsic properties of the ODE learned rectified flow, and can achieve one-step simulation for most of the states, making it much faster for eiciisionmaking tasks in real-world applications. Recently, new generative models [2 25 22 23 46 24 26 47] have emerged. These models directly learn probability flow ODEs by constructing linear interpolations between two distributions, or learn to distill a pretrained diffusion model (26 47 with an additional distillation training phase. Empirically, these methods exhibit more efficient inference due their preference of straight trajectories. Among them, Rectified flow achieves one step generation with reflow, process to straighten the ODE. However, it requires a costly synthetic data construction.
By contrast, AdaFlow only leverages the initially learned ODE, but still keeps cheap training and inference costs that are similar to behavior cloning. We achieve this by unveiling a previously overlooked feature of these flow-based generative models: they act as one-step generators for deterministic target distributions, and their variance indicates the straightness of the probability flows for a certain state. Leveraging this feature, we design AdaFlow to automatically change the level of action modalities given on the states.
Diffusion Models for Decision Making For decision making, diffusion models obtain success as in other applications areas 51) In a pioneering work, Janner et al. [3] proposed "Diffuser",
planning algorithm with diffusion models for offline reinforcement learning. This framework
extended to other tasks in the context of offline reinforcement learning [52], where the training
dataset includes reward values. For example, Ajay et al. [14] propose model policies as conditional diffusion models. The application of DDPM [16] and DDIM [3] on visuomotor policy learning for physical robots [6] outperforms counterparts like Behavioral Cloning. Freund et al. [53] exploits two coupled normalizing flows to learn the distribution of expert states, and use that as a reward to train an RL agent for imitation learning. AdaFlow admits much simpler training and inference pipeline compared with it. Despite the success of adopting generative diffusion models as decision makers in previous works, they also bring redundant computation, limiting their application in real-time, low-latency decision-making scenarios for autonomous robots. AdaFlow propose to leverage rectified flow instead of diffusion models, facilitating adaptive decision making for different states while significantly reducing computational requirements. In this work, similar to Diffusion Policy [6] we focus on offline imitation learning. While AdaFlow could theoretically be adapted for offline reinforcement learning, we leave it for future works.
3 AdaFlow for Imitation Learning
To yield an agent that enjoys both multi-modal decision-making and fast execution, we pro- pose AdaFlow, an imitation learning framework based on flow-based generative policy. The merits of AdaFlow lie in its adaptive ability: it identifies the behavioral complexity at state before allocating
3

---

computation. If the state has deterministic choice of action, it outputs the required action rapidly; otherwise, it spends more inference time to take full advantage of the flow-based generative policy. This handy adaptivity made possible by leveraging combination of elements: 1) a special property of the flow 2) variance estimation neural network and 3) a variance-adaptive ODE solver. We formally introduce the whole framework in the sequel.
3.1 Flow-Based Generative Policy
Given the expert dataset D {((() a ) trajectories following the target distribution TE. The can be induced from a state-conditioned flow- based model,
i=1' goal is to learn policy To that can generate
dz = ( s) 2o ~ NN.
Here, $ is the state and the velocity field is parameterized by neural network 0 that takes the state as an additional input. To capture the expert distribution with the flow-based model, the velocity field can be trained by minimizing state-conditioned least-squares objective,
I(O)=
E (s.ajeD Mo ry
||a xo vol t. {))1 1012 dt
where zy is the linear interpolation between 2g and 21 = a:
ta + (1 t)
(3)
We should differentiate 2, which is the ODE trajectory in ① from the linear interpolation æt. They do not overlap unless all trajectories of ODE ① are straight. See Liu et al. [21] for more discussion.
With infinite data sampled from NE. unlimited model capacity and perfect optimization, it is guaran- teed that the policy me generated from the learned flow matches the expert policy TE (21].
3.2 The Variance-Adaptive Nature of Flow
Typically, to sample from the distribution Fe at state we start with a random sample zo from the Gaussian distribution and simulate the ODE (Eq. (I) with multi-step ODE solvers to get the action. For example, we can exploit N-step Euler discretization,
N % (Z...t) 8), f, N Osi<
(4)
After running the solver, 21 is the generated action. This solver requires inference with the network N times for decision making in every state. Moreover, a large N is needed to obtain a smaller numerical error.
However, different states may have different levels of difficulty in deciding the potential actions. For instance, when traveling from city A to another city B, there could be multiple ways for transportation, corresponding to a multi-modal distribution of actions. After the way of transportation chosen, the subsequent actions to take will be almost deterministic. This renders using uniform Euler solver with the same number of inference steps N across all the states a sub-optimal solution. Rather, preferred that the agent can vary its decision- makiin process as the state of the agent changes. The challenge is how to quantitatively estimate the complexity of state and employ the estimation to adjust the inference of the flow-based policy.
Variance as a Complexity Indicator We notice an intriguing property of the policy learned by rectified flow, which connects the complexity of state with the training loss of the flow-based policy: if the distribution of actions is deterministic at a state 8 (i.e., a Dirac distribution). the trajectory of rectified flow ODE is straight line, i.e., a single Euler step yields an exact estimation of 21.
Proposition 3.1. Let o² be the optimum of Eq. ② var (a s) = 0 where a Tg(- s), then the learned ODE conditioned on S is
dz, = v² (Zz, 1 8)dl = (a (a-20)d4 20), VL E (o, 1).
(5)
whose trajectories are straight lines pointing to and hence can be calculated exactly with one step of Fuler step:
ø1 20 + 2² (z0. 0ls)

---

Algorithm AdaFlow: Execution
1: Input: Current state 8, minimal step size emin, error threshold 7/, pre-trained networks to and de-
2: Initialize 2 ~ 00, 1
3: while do
4: Compute step size
6: = Clip c (8)
17
[Emino - o1
5: Update t + c2, + 21 + c (
6: end while
7: Execute action 21.
Note that the straight trajectories of satisfies z, = ta + (1 z0. which makes it coincides with the linear interpolation 24,- As show 21. this happens only when all the linear trajectories do not intersect on time E 1).
More generally, we expect that the straightness of the ODE trajectories depends how determin- istic the expert policy T is. Moreover, the straightness can be quantified by conditional variance metric defined as follows:
o² (zx, var (a
Io x 2, s)
E
||a
Za - v² (xe,
s)| x =
Proposition 3.2. Under the same condition as Proposition 3.1 we have o² o² t| = 0 from (5)
The proof of the above propositions is in Appendix A. To summarize, the variance of the state- conditioned loss function at (z,, t) can be an indicator of the multi-modality of actions. When the variance is zero, the flow-based policy can generate the expected action with only one query of the velocity field, saving huge amount of computation. In Section 3.3 we will show the variance can be used bound the discretization error, thereby enabling the design of an adaptive ODE solver.
Variance Estimation Network In practice, the conditional variance o² s) can be empirically approximated by neural network of (a, s) with parameter Φ. Once the neural velocity vo is learned, estimate σg by minimizing the following Gaussian negative fog-likelihood loss:
min EE J. la Dg - ug (z ||) + log (2), ([s)d(
(x, ()))
We adopt two-stage training strategy by first training the velocity network Ug then the variance estimation network og. In practice, the second stage just involves fine-tuning few linear layers on top of the trained velocity network. Alternatively, we can optimize both the variance estimation and action generation simultaneously, which can extend training time. Our experiments show that joint training and two-stage training yield comparable performance.
3.3 Variance-Adaptivv Flow-Based Policy
Because the variance indicates the straightness of the ODE trajectory, it allows us develop adaptive approach to set the step size to yield better estimation with lower error during inference.
To derive method, let us consider to advance the ODE with step size e, 2,:
Zi+c, = 2t + cg2" (2t, (s).
(8)
The problem is how to set the step size e properly. If c, is too large, the discretized solution will significantly diverge from the continuous solution; if e, is too small, will take excessively many steps to compute.
propose an adaptive ODE solver based on the principle of matching the discretized marginal distribution Pt of from (8). and the ideal marginal distribution P when following the exact ODE
5

---

1 This is made possible with key insight below showing that the discretization can be bounded by the conditional variance o2(
Proposition 3.3. Let pi be the marginal distribution of the exact ODE dz, = vz,, s)dt. Assume
2e Pt Př and pite, the distribution of Zite, following (8) Then we have
We Ptte.) Ezipp [o (ze,t |s]]
where W2 denotes the 2- 2-Wasserstein distance.
We provide the proof in Appendix A.2
Hence, given a threshold 77, to ensure that an error of (2( P n/². we can bound the step size by e S n// s). Because e, time should not be large than 1 - t, we suggest the following rule for setting the step size at 2, at time t:
= Clip Clip
"
[Cmin: 1 1-
(9)
( s)²
where we impose an additional lower bound min to avoid 6z to be unnecessarily small. Besides, the proposed adaptive strategy guarantees to instantly arrive at the terminal point when o² (z,, s) as 6 = 1 - f. Moreover, it aligns with Section. 3.2 since for states with deterministic actions, sets 6D 1 to generate the action in one step. We incorporate the above insights to the execution in Algorithm
Global Error Analysis Proposition 3.3 provides the local error at each Fuler step. In the following, we provide an analysis of the overall error for generating when we simulate ODE while following the adaptive rule (9). To simplify the notation, we drop the dependency on the state 8, and write o7 (-) = (,,t | s).
Assumption 3.4. Assume ||u] Lip < L for E [0, and the solutions of dz = vz(2)) has bounded second curvature | < for E [0.
This is a standard assumption numerical analysis, under which Euler's method with a constant step size of Emin admits global error of order ((Emia).
Proposition 3.5. Under Assumption 3.4 assume follow Euler step (8) with step size er in (9)
starting from Zg = 3B) ~ Py- Let Pt be the distribation of zt. we obtained in this way, and Pt that
of az in (3) Note that pi is the true data distribution. Set vi Mi 2 for some Mn > 0, and Emin 1/Nmax-
Let Nads be the number of steps we arrive at 21 following the adaptive schedule. We have
N Nada
Wa(pi.P1) S
X émins
N Nmax
where C' is constant depending M, My and L.
The idea that the error is proportional to Nada suggesting that the algorithm claims an improved Near bound in the good case when it takes smaller number of steps than the standard Fuler method with constant step size emin- We provide the proof in Appendix A.3
Discussion of AdaFlow and Rectified Flow. Rectified Flow operates in two stages: the first is learning an ordinary differential equation (ODE), and the second involves a technique called "reflow" used to straighten the learned trajectory. Theoretically, reflow allows for one-step action generation. However, using reflow introduces two major drawbacks 1) It significantly prolongs training time, particularly because generating the required pseudo noise-data pairs through ODE simulation is computationally expensive; 2) It leads to poorer generation quality due to straightened ODE. In contrast, method utilizes only the original ODE, eliminating the need for additional reflow or distillation process, and consistently achieves more accurate action generation.
4 Experiments
We conducted comprehensive experiments on four sets of tasks: 1) a simple 1D toy example to demonstrate the computational adaptivity of AdaFlow; 2) 2D navigation problem; and two robot

---

manipulation task suites 3) RoboMimic [14 following past works 6 and 4) LIBERO diverse and realistic scenarios for evaluation.
provide
Our results show that AdaFlow improves the success rate of completing both navigation and manipu- lation tasks, outperforming state-of-the-arl methods such as BC and its variants, as well as Diffusion Policy, across range of tasks. Additionally, AdaFlow drastically reduces the inference cost. Further experiments demonstrate that AdaFlow is robust to changes in hyperparameters and can adaptively adjust its inference speed according to different states, ensuring efficient and reliable performance.
BC
Diffusion Policy
Rectified Flow
AdaFlow
Behavior Diversity
x /
Fast Action Generation
/
X
No Distillation Rellow
Table I: Comparison between BC, Diffusion Policy, Rectified Flow and AdaFlow.
4.1 Regression
We start with ID regression task designed demonstrate the adaptivity nature of AdaFlow. The goal is to learn a mapping from æ to where
for æ S 0
(10)
- for x >
Note that y is deterministic when x S O and stochastic otherwise. The training and testing data are uniformly sampled from the ground-truth function with æ E [- 5]. Details about the setup and the hyperparameters are provided Appendix.
AdaFlow can achieve 1-step generation for deterministic states. Figure (top-right) shows the generation trajectories of Diffusion Policy and AdaFllow with step. Notably, when x S AdaFlow generates straight trajectories and is therefore able predict with a single step, aligning our analysis Proposition 3.1 and 3.2 In contrast, Diffusion Policy generates curved trajectories when step = and hence cannot predict accurately with a single step. The bottom of Figure 2]shows the estimated variance by AdaFlow across E [-5, 5]. which accurately aligns with the expected variance. In addition, as z increases, AdaFlow adaptively increases the required number of simulation steps.
4.2 Navigating a 2D Maze
We create sets of maze navigation tasks validate AdaFlow's performance of modeling multi- modal behavior. In particular, we create two single-task environments where the agent starts and ends fixed point and two multi-task environments where the agent can start and end at different points. All four environments are simulated in D4RL Maze2D 54 using MuJoCo. The environments and demonstrations are visualized in Figure 7
Diffit
Difl-n (NFE-20)
BC (NFE-1)
AdaFlew (NFE=1.12)
Method
NFE, Maze
Maze Maze Maze
Needs reflow
Rectified Flow
0.82
1.00
1.00
0.80
BC
1.00
1.00
0.92
0.76
BC-GMM
0.84
1.00
0.88
0.72
Diffussion Policy
0.00
0.32
0.16
0.08
Diffussion Policy
5
0.58
1.00
0.84
0.76
Diffussion Policy
20
0.62
0.98
0.84
0.82
AdaFlow
1.56
0.98
1.00
0.96
0.86
Table 2: Performance on maze navigation tasks. The table showcases success rate each model across different maze complexities. The highest success rate for each task are highlighted in bold. NFE denotes Number of Function Evaluations.
Figure 3: Generated trajectories. visualize trajectories generated by different policies, with the agent's starting point fixed.

---

Lift
Can
Square
Transport
ToolHang
Push-T
Method
NFE
ph
mh
ph
mh
ph
mh
ph
mh
ph
ph
Rectified Flow (Needs reflow)
1.00
1.00
0.94
1.00
0.94
0.92
0.90
0.76
0.88
0.92
LSTM-GMM
1.00
1,00
1.00
1.00
0.95
0.86
0.76
0.62
0.67
0.69
IBC
0.79
0.15
0.00
0.01
0.00
0.00
0.00
0.00
0.00
0.75
BET
1.00
1.00
1.00
1.00
0.76
0.68
0.38
0.21
0.58
Diffusion Policy
0.04
0.04
0.00
0.00
0.00
0.04
Diffusion Policy
0.64
0.98
0.52
0.66
0.56
0.12
0.84
0.68
0.68
0.34
Diffusion Policy
100
1.00
1,00
1.00
1.00
1.00
0.97
0.90
0.72
0.90
0.91
AdaFlow
1.17
1.00
1.00
1.00
0.96
0.98
0.96
0.92
0.80
0.88
0.96
Table 3: Success rate on RoboMimic Benchmark. The highest success rate for each task are highlighted in bold.
Open top drnwer of the cabinet
black the back on the plase
Put the black bowl at the front on the plate
Put the middle black bosl on the plate
Put the middle black bowl the cabinet
Stack front black the black bowl in the middle
Task
Task 2
Task
Task 4
Task 5
Task
Start
End
Figure 4: LIBERO tasks. We visualize the demonstrated trajectories of the robot's end effector.
AdaFlow achieves high diversity and success with low NFE; Diffusion Policy and BC lag in comparison. We compare AdaFlow against baseline methods in Table We additionally visualize the rollout trajectories from each learned policy in Figure 33 as qualitative comparison of the learned behavior across different methods. From the results, we see that AdaFlow with an average Number of Function Evaluation (NFE) of 1.56 NFE achieve highly diverse behavior and high success rate in the meantime. By contrast, Diffusion Policy only demonstrates diverse behavior when NFE is larger than 5 and falls behind in success rate even with 20 NFE compared AdaFlow. BC, on the other hand, has high success rate while performing relatively poorly in terms of behavior diversity.
4.3 Robot Manipulation Tasks
Experiment Setup To further validate how AdaFlow performs practical robotics tasks, compare AdaFllow against baselines on Push-T task (6), the RoboMimic [10] benchmark (Lift, Can, Square. Transport, ToolHang) and the LIBERO (7 benchmark. For the Push-T task and the tasks in RoboMimic, we follow the exact experimental setup described in Diffusion Policy [6]. Following the Diffusion Policy, add three additional baseline methods: 1) LSTM-GMM, BC with the LSTM model and Gaussian mixture head, 2) IBC, the implicit behavioral cloning [2), an energy-based model for generative decision-male of six Kitchen tasks and follow the setup described in the LIBERO paper (Check Figure 4 for the description of the six tasks).
and 3) BET 11 For the LIBERO tasks, we pick subset
Method
NFF.
Task
Task 2
Task 3
Task
Task
Task 6
Average
Rectified Flow (Needs reflow)
0.90
0.82
0.98
0.82
0.82
0.96
0.88
Diffusion Policy
0.00
0.00
0.00
0.00
0.00
Diffusion Policy
2
0.58
0.36
0.66
0.36
0.32
0.38
Diffusion Policy
20
0.94
0.84
0.98
0.78
0.82
0.92
0.88
AdaFlow
1.27
0.98
0.80
0.98
0.82
0.90
0.96
0.91
Table 4: Success Rate on LIBERO Benchmark. The highest success rate for each task are highlighted in bold.
AdaFlow consistently outperforms competitors varied robot manipulation tasks with high efficiency. The results of the Push-T task and the RoboMimic benchmark are summarized in Table 3 From the table, we observe that AdaFlow consistently achieves comparable higher success rates across different challenging manipulation tasks, compared against all baselines, with only average

---

NFE of 1.17. Note that Diffusion Policy, while showing high success rates using NFE 100, falls behind when NFE = 1. Results for the six LIBERO tasks are presented in Table 4 Aligning with findings from our previous experiments. AdaFlow once again outperforms BC and Diffusion Policy in terms of success rate with an average NFE of .277 We additionally visualize the variance predicted by AdaFlow Figure It can be seen that the model identifies the high variance when the robot's end-effector is close to the object or target area, matching the variance from the demonstration data.
4.4 Ablation Study
We valid how AdaFlow performs against baselines regarding the training and inference efficiency, In addition, we examine how critical the variance estimation network is.
Ground Truth Variance
BC
Diffusion Policy
AdaFlow (Ours)
1.00
0.75
X
0.25
Predicted Variance
0.00
20
NFB
1.00
0.75
& 0.50
Epoch
Figure Predicted variance. We visualize the variance predicted by AdaFlow. The variance computed states from the expert's demonstration and eraged over simulation steps (c.g., from to 1). Then we normalize the variance to [0, 1] by the largest variance found at all states.
Figure 6: Ablation studies on AdaFlow.
Higher Training and Inference Efficiency. Figure examines changes in success rate relative to the NFE. AdaFlow maintains high success rate with a very low NFE, whereas the Diffusion Policy generally requires more than three NFE to perform well. Although BC performs well with one NFE, it demonstrates very limited behavioral diversity and struggles to model multi-modal behavior. Figure 6 (bottom) illustrates training efficiency by displaying the success rate epochs. It shows that AdaFlow has a better area-under-cure than Diffusion Policy, indicating faster learning. As expected, due to its simplicity. Behavioral Cloning (BC) achieves the best learning efficiency.
Robustness to 7) In Figure 6 the NFEs in AdaFlow are calculated at various values, shows that AdaFlow is robust to changes in
On the Importance of Variance Estimation. Table provide the performance of AdaFlow with and without the variance estimation network on the four mazes from Section 4.2] From the results, is clear that the variance estimation network not only makes inference faster, but can also lead to better performance.
Maze1
Maze1
Maze3
Maze4
w/o Variance Estimation AdaFlow (Ours)
0.78
1.00
0.92
0.80
0.98
1.00
0.96
0.86
Table Ablation study the use of estimated variance to determine inference steps. Buler sampler used for AdaFlow without variance estimation.

---

5 Conclusion
We present AdaFlow, novel imitation learning algorithm adept at efficiently generating diverse and adaptive policies, addressing the trade-off between computational efficiency and behavioral diversity inherent in current models. Through extensive experimentation across various settings, AdaFlow demonstrated superior performance across multiple dimensions including success rate, behavioral diversity, and traiiin//xxecutin efficiency. This work lays robust foundation for future research on adaptive imitation learning methods in real-world scenarios.
References
[1] Schaal, S. Is imitation learning the route humanoid robots? Trends in cognitive sciences, 3(6):233-242, 1999.
[2] Osa, T., J. Pajarinen, G. Neumann, et al. An algorithmic perspective on imitation learning. Foundations and Trends in Robotics, 7(1-2)11-179 2018.
[3] Liu, B.. X. Xiao, P. Stone. lifelong learning approach to mobile robot navigation. lEEE Robotics and Automation Letters, 6(2):1090-109)) 2021.
[4] Zhu, Y., A. Joshi, P. Stone, et al. Viola: Imitation learning for vision-based manipulation with object proposal priors. 6th Annual Conference Robot Learning (CoRL), 2022.
[5] Brohan, A., N. Brown, J. Carbajal, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint wrxXiv22222068111 2022.
[6] Chi, C., S. Feng, Y. Du, et al. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
[7] Pomerleau, D. A. Alvinn: An autonomous land vehicle in neural network. Advances in neural information processing systems, 1, 1988.
[8] Ross, S., G. Gordon, D. Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statisties, pages 627-635. JMLR Workshop and Conference Proceedings. 2011.
[9] Torabi, F., G. Warnell, P. Stone. Behavioral cloning from observation. In Proceedings of the Twenty-Sevnnth International Join! Conference on Artificial Intelligence. International Joint Conferences on Artificial Intelligence Organization, 2018.
[10] Mandlekar, A., D. Xu, J. Wong, et al. What matters learning from offline human demonstra- tions for robot manipulation. arXiv preprint arXiv:2108.0322 2021.
[11] Shafiullah, M., Z. Cui, A. A. Altanzaya, et al. Behavior transformers: Cloning & modes with one stone. Advances in neural information processing systems, 35:22955-22963. 2022.
[12] Florence, P., C. Lynch, A. Zeng, et al. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168. PMLR, 2022.
[13] Janner, M., Y. Du, Tenenbaum, et al. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pages 9902-9915. PMLR, 2022.
[14] Ajay, A., Y. Du, A. Gupta, et al. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2022.
[15] Song, Y.. J. Sohl-Dickstein, D. P. Kingma, et al. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations. 2020.
[16] Ho, J., A. Jain, Abbeel. Denoising diffusion probabilistic models. Advances neural information processing systems, 33:6840-6851, 2020.
10

---

[17] Sridhar, A., D. Shah, C. Glossop, et al. Nomad: Goal masked diffusion policies for navigation and exploration. arXiv preprint arXiv:2310.07880 2023.
[18] Team, O. M., D. Ghosh, H. Walke, et al. Octo: An open-source generalist robot policy, 2023.
[19] Shah, D., A. Sridhar, N. Dashora, al. Vint: A foundation model for visual navigation. arXiv preprini arXiv:2306. 14846, 2023.
[20] Hansen-Estruch, P., I. Kostrikov, M. Janner, et al. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
[21] Liu, X., C. Gong, Q. Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations. 2022.
[22] Lipman, Y., R. T. Chen, H. Ben-Hamu, al. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations. 2022.
[23] Albergo, M. S., N. M. Boffi, E. Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiiv23333007977, 2023.
[24] Heitz, E., L. Belcour, T. Chambon. Iterative -dde)blending: A minimalist deterministic diffu- sion model. In ACM SIGGRAPH 2023 Conference Proceedings, SIGGRAPH '23. Association for Computing Machinery, New York, NY, USA, 2023.
[25] Liu, Q. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 2022.
[26] Song, Y., P. Dhariwal, M. Chen, et al. Consistency models. 2023.
[27] Liu, B., Y. Zhu, C. Gao, al. Libero: Benchmarking knowledge transfer for lifelong robot learning, arXiv preprint arXiv:2306. 03310, 2023.
[28] Sohl-Dickstein, J., E. Weiss, N. Maheswaranathan. et al. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[29] Song, Y. S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[30] Ho, J.. T. Salimans, A. A. Gritsenko, et al. Video diffusion models. In Advances in Neural Information Processing Systems.
[31] Zhang, S., X. Yang. Y. Feng, et al. Hive: Harnessing human feedback for instructional visual editing. arXiv preprint arXiv:2303.096 09618, 2023.
[32] Wu, L.,, C. Gong, X. Liu, et al. Diffusion-based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems, 25:36533-36545 2022.
[33] Saharia, C., W. Chan, S. Saxena, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36499 2022.
[34] Kong, Z., W. Ping, J. Huang, et al. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations. 2020.
[35] Luo, S., W. Hu. Diffusion probabilistic models for 3d point cloud generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2837-2845. 2021.
[36] Liu, L. Wu, M. Ye, et al. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations. 2023.
[37] Luo, S., W. Hu. Score-based point cloud denoising. Proceedings of the IEFF/CVF Interna- tional Conference on Computer Vision, pages 4583-4592, 2021.

---

[38] Wu, I.., D. Wang, C. Gong, al. Fast point cloud generation with straight flows. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 9445-9455 2023.
[39] Karras, T., M. Aittala, T. Aila, et al. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems. 2022.
[40] Lu, C., Y. Zhou, F. Bao, et al. 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.
[41] Liu, L., Y. Ren, Z. Lin, et al. Pseudo numerical methods for diffusion models manifolds. In International Conference on Learning Representations. 2021.
[42] Lu, C., Y. Zhou, F. Bao, et al. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095, 2022.
[43] Song, J., C. Meng. S. Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations. 2020.
[44] Bao, F,, C. L. Zhu, et al. Analytic-dpm: analytic estimate of the optimal reverse variance in diffusion probabilistic models. In International Conference on Learning Representations. 2021.
[45] Press, W. H., S. Teukolsky. Adaptive stepsize runge-kutta integration. Computers in Physics, 6(2): 188-111, 1992.
[46] Albergo, M. S., E. Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations. 2022.
[47] Salimans, T., J. Ho. Progressive distillation for fast sampling of diffusion models. In Interna- tional Conference on Learning Representations, 2022.
[48] Kapelyukh, I., V. Vosylius, E. Johns. Dall-e-bot: Introducing web-scale diffusion models to robotics. IEEE Robotics and Automation Letters, 2023.
[49] Yang, S., O. Nachum, Y. Du, et al. Foundation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.041299 2023.
[50] Pearce, T., T. Rashid, A. Kanervisto, et al. Imitating human behaviour with diffusion models. In Deep Reinforcement Learning Workshop NeurIPS 2022. 2022.
[51] Chang, W.-D.,
J. C. G. Higuera, S. Fujimoto, al. 1l-flow: Imitation learning from observation using normalizing flows. arXiv preprint arXiv:2205.09255, 2022.
[52] Wang, Z., H. Zheng, P. He, et al. Diffusion-gan: Training gans with diffusion. In The Eleventh International Conference on Leurning Representations. 2022.
[53] Freund, J., E. Sarafian, S. Kraus. A coupled flow approach to imitation learning. In International Conference on Machine Leurning, pages 10357-10372. PMLR, 2023.
[54] Fu, J., A. Kumar, O. Nachum, et al. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219. 2020.
12

---

A Appendix
A.1 Proof of Proposition 3.1 and Proposition 3.2
PROOF 1. varzz (a]s) = O means that the action = a equals a deterministic value given
With 3; = ta + (1 t) To, note that
DD (a Tz)-
Therefore, a zo is deterministically decided by æe and s. This yields
(x, 8) Ela Zp a. = x)
(a x).
Therefore, we have
dzg = (2t, 11)) 1 (a - 21)
Solving ODE this yields
2t to + (1- f)== (1 t)y" (zo, 0 01s).
Differentiating it also yields
2: = (a nn dt..
We also have (z 8) again because
is deterministic given and
o² 0² var(a Wo ar x. 8) 0.
A.2 Proof of Proposition 3.3
PROOF Following the property of rectified flow, the distribution of x1 + (1 t) coincides with Pt for all I. E 10, 1). Hence, can assume that 2t. It pt. In this case, we have Zeta = 2z + 4,00 (zr, 8) and aite, = 24 + ex(a zo). We have
W2 W2(pite, )2tee²
S E
|ztte Ieteilll
E ||zttss 21teell2 Zz
E
E E |lexv* (ze, () -ei(a - (c(a-=zo)13 at
Ez ( |)]]
A.3 Proof of Proposition 3.5
PROOF Assume the adaptive algorithm visits the time grid of 0 to, ty 1.
Define z² be the result when we implement the adaptive discretization algorithm upto t; and then switch to follow the exact ODE afterward, that is, we have dz; = ve(2 dt for t 2 t4. In this way, we have 2 2t, and z z;, where 2t is the trajectory the exact ODE dæ; ut (z7 dt.
From Lemma
we have
2 24-1
14-1-2i1 2 exp([(11
2//
2
Let p be the distribution of zt. Then we have pt = and pt = př. Then
1/2
W2(1--' ptt < E
2
4:-1 2t
1 2/: 17.1131
1/2
= exp(T.(1 - 1:))
exp(L(1 - (:)) max(n, 2 M/2)
Cea. (
13

---

where C 1 max M, Mn) exp(I(( 14)). Here the bound in the proof Proposition 3.1
and Lemma A. Hence,
Nada
W2(Pi,P1) = Σ W2(1i ,2f)
==1
Mado
Cemin exp( (Lt,)
Nada
$ C
X Cmint
Nmax
where C = eXp(L) max M,, M,).
Lemma A.1. Let ||o. Lop I L for E [0, 1]. Assume da and y solve dæ, ve(xs)dt and dy: = ve(y))dt starting from no. yo- respectively. We have
|o well I exp(Lt) 1|000 - 3oll Vt E [0,1].
(11)
PROOF
| yall² = 2(xx y1) TT)TT(((()))) (ux(xe) - v. (1n))
dt
S |00 vell?
where we used (xt) - (Be)ll < L || - y Using Gronwall's 's inequality yields the result.
Lemma A.2. Under Assumption 3.4 we have
2
||aete - (xe + cvx(
2
for < I1še+1S1.
PROOFS 5. Direct application of Taylor approximation.
A.4
Visualization of Tasks
We provide visualization of the Maze Figure] 7
Start
Goal
Demonstration (Single-task)
Demonstration (Multi-task)
Maze
Maze 2
Maze 3
Maze 4
Figure 7: Trajectories of 100 demonstrations for each maze.
A.5 Planner for Maze2D task
We generate the demonstration data in Maze toy using planner similar [44 The planner devises path in a maze environment by calculating waypoints between the start and target points. It begins by transforming the given continuous state space into discretized grid representation. Employing Qelannnng, it evaluates the optimal actions and subsequently computes the waypoints by performing a rollout in the grid, introducing random perturbations to the waypoints for diversity. The controller connects these waypoints in ordered manner to form feasible path. In runtime, it dynamically adjusts the control action based on the proximity to the next waypoint and switches waypoints when close enough, ensuring the trajectory remains adaptive and efficient.
14

---

Figure 8: Predicted variance by AdaFlow on the Maze task.
A.6 Comparative Analysis of Separate and Joint Training
In this section, we provide a comparison between the two training strategies employed in our proposed solution: separate training and joint training. Our primary objective is investigate whether there is a substantial difference in performance and efficiency between these two training approaches.
Experiment Setup. To conduct this comparative analysis, we designed experiments using our proposed framework with both training strategies. Specifically, we consider two approaches: separate and joint training. In Separate Training setting, we train the variance prediction network and the policy function separately, as described our main paper. In Joint Training setting, we train both the variance prediction network and the policy function simultancously in end-to-end manner. The goal is to assess the impact of these training strategies on the overall performance.
Results and Discussion. As shown in Table 6 the performance were consistent between the two training approaches, indicating the effectiveness our two-stage framework in balancing policy accuracy and uncertainty estimation. Separate training exhibited faster computational speed, making the preferred choice once the policy function was robustly trained. Joint training required more computational resources and time.
Maze
Maze 2
Maze 3
Maze 4
AdaFlow (Separate)
0.98
1.00
0.96
0.86
AdaFlow (Joint)
1.00
1.00
0.96
Table 6: Performance comparison of separate training and joint training of AdaFlow in Maze tasks.
A.7 Visualization Exact Variance.
In the main paper, showed the variance predictions made by AdaFlow across different states within a robot's state space. Here, we explain how we compute the exact variance for different states, to provide a ground truth of variance for reference. To achieve this, first train 1-Rectified Flow model for the task, then we can compute the exact variance by sampling:
1 ΣΣ zu u(z,, (2)]] where 2, ty + (1 t) (x, v) ~ pp. (12)
N₂ N
=p
For each states, we randomly sample 10 time steps (N₃ 10) and 10 noises (N₃ 10).
A.8
Visualization of Predicted Variance on Maze task.
We present the predicted variance by AdaFlow in Figure
A.9
Additional Experimental Details.
Model Architectures. For the ID toy example, we used a MLP constructed with fully connected layers and SiLU activation functions. We integrated temporal information by extending the time input into a 100-dimensional time-encoding vector through the cosine transformation of a random vector, cost 27, where is sampled from Gaussian distribution. This time feature is then concatenated
15

---

with the noise and condition (α) inputs to for time-aware predictions. The network comprises 4 hidden layers, each with 100 neurons, and the output layer predict a single y value. The dataset consists of 10000 single-dimensional samples uniformly distributed in the range [-5,5].
For navigation and robot manipulation tasks, we adopted the model architecture from Diffusion Policy [6]. For navigation task, we use the same architecture used in Push-T task. In the RoboMimic and LIBERO experiments, we used the Diffusion Policy-CC architecture. To ensure a fair comparison across different methods, we maintained consistent architecture for all methods in experiments, except where specifically noted. Detailed parameters are available in Table 7
Hyperaommrr
RF. AdaFline
TNC
TXffbckir Fallcy
RF ArB
IKC
Diiffsiin Polley
RIP Adaffive
BC
Diffecion Priicy
Optinier
Adan
AdooW
Aaann
AlleW
Adant
Aduus
0.9
0.9
0.9
0.9
0.4
0.5
0.95
0:95
D.55
D/599
0.999
13.099
0.550
0.999
0.550
0/99
1.590
0.999
Barc
1000
1000
1000
256
256
64
64
Tipnchs
203
200
170
200
2001
300
500/ 1/ TDXKRRM)
SO0IE 3000(RM)
S0O0 3000RR)
Leeiing nate schedater
COCT0
coulee
Onliee
costee
coder
coene
costne
coeite
coder
Narbes of training litee stogs Surber of ficerowee Sime stes
100
20
100
100 RE))
DONDDPMI
20DDEPMI
1ONDEPPN)
Actios prodiction locavo
16
16
slecrvation input Aatis esecation
Olograaion sizz
(Siegll tack) Mo
76 70
TOx 76
70 X.
Table Hyperparameters used for training AdaFlow and baseline models.
Implementation of Baselines. In our studies, BC was implemented as a baseline, applying behavior cloning in its most straightforward form and using a Mean Squared Error loss function between the predicted and ground truth actions. The implementations for DDPM and DDIM remained consistent with those outlined in (6). Across all experiments, consistency was maintained regarding architecture. input, and output, with all methods adhering to similar experimental pipeline. We just use a 4 layer MLP with SiLU activation for the variance prediction, with hidden dimension of 512, which is very small network whose computational overhead can be neglected compared to the full model.
Implementation of Vairance Prediction Network. In the ID toy experiment, we designed the variance prediction network as a 4-layer MLP, mirroring the main model's architecture for simplicity. In theory. the variance estimation network takes the same input as rectified flow model, its input can be just the intermediate features extracted by the main model. Hence in the navigation and manipulation experiments, the inputs of variance prediction networks are the bottle-neck features extracted by the U-Net model.
Training on Robo.innic. Training Diffusion Models on RoboMimic is very resource-intensive. Training and evaluating a Transport task requires over a month of GPU hours. More complex tasks, such as ToolHang, can demand up three times longer [Gvven the challenges in replicating the results from (6. we opted to start with their open-sourced pretrained model. We then fine-tuned the baselines and our method for 500 epochs and subsequently compared the performance of different models.
A.10 Comparison with standard Rectified Flow.
For the purpose of policy learning, can consider standard Rectified Flow as subset of our method, which can be recovered with specific choices of 1) and émin In this section, we compare our approach with the standard Rectified Flow, particularly focusing on the generation within a single step. Standard Rectified Flow requires a reflow or distillation stage to straighten the ODE process. During this reflow stage. the model simulates data using the initial 1-Rectified Flow. These data are then used in distillation training, resulting in what is termed a 2-Rectified Flow. Theoretically, 2-Rectified Flow is capable of producing straight generation trajectory, which enables one-step generation. In contrast, the 1-Rectified Flow tends to be less straight, necessitating multiple steps for sample generation.
In T we compare the performance I-Rectified Flow, 2-Rectified Flow, and our method the maze task. Furthermore, Figure iillustrates the trajectories produced by both standard Rectified Flow
'See this link
16

---

1-Rectified Flow (NFE=1)
1-Rectified Flow (NFE-5)
2-Rectified Flow (NFE=1)
AdaFlow (NFE=1.12)
Figure 9: Generated trajectories. visualize trajectories generated by standard Rectified Flow and AdaFlow, with the agent's starting point remaining fixed. 0
and our method. It's evident that the standard I-Rectified Flow struggles to generate a diverse range of actions in single step. In contrast, our method is able produce diverse behaviors in nearly one step. This efficiency attributed to our method's s ability to estimate the variance across different states, identifying those that require multi-step generation.
NFEL
Maze
Maze 2
Maze
Maze
I-RF
1
1.00
1.00
0.98
0.80
I-RF
5
0.82
1.00
0.94
0.80
2-RF (reflow)
0.82
1.00
AdaFlow (n 1.5)
1.56
0.98
1.00
0.96
0.86
AdaHlow (1 2.5)
1.12
1.00
1.00
0.78
Table 8: Performance maze navigation tasks. table showeases success rate (SR) for each model across different maze complexities. The highest success rate for each task are highlighted in bold. Note that 2 2- needs an expensive distillation training stage.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
18

---

Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced
statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
19

---

Answer: [Yes]
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (kttps://nnis.co cc/
mublic/guides////
Soardd/s//ooeSStmmmiiiiooooooo. Policcyy)))) for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https for more details.
/ //nip. https/cccpbbllccuuuccceeee wttpiii/ggii///e/eeeee getggdddde//eeubbmmiiimmcc// seriiisssinpooiiio
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
20

---

It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: htte:://neuuiie nwwrips....../. httpp/bccc/thhiccccccee.........
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
21

---

generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi www.whethco...co.. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
22

---

Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codc/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions NeurlPS Code of Ethics and the
and locations, and we expect authors to adhere
guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
23


Reviews:
[{"review_id": "2f6GxQtuTd", "reviewer": "Reviewer_df3a", "summary": "This paper proposes the AdaFlow algorithm for training fast and accurate imitation learning agents. AdaFlow learns a policy based on the flow generative model and makes use of the connection between the conditional variance of the training loss and the discretization error of the ODEs. In practice, it uses the estimated conditional variance to dynamically adapt the step size of the inference model. By doing so, it reduces inference time compared to diffusion models, and as the experiments show, learns a more accurate model.", "questions": "* In line 174 - should the inequality for epsilon be with $\\sqrt{\\eta}$ instead of $\\eta$?\n* Can the adaptive step be combined with a better ODE solver, such as RK4 (https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods)? if not, does AdaFlow perform better than RK4?\n* Did you check other methods to estimate the uncertainty in action prediction (other than training a network that estimates the conditional variance)? For example, one can train an ensemble of agents and make use of the variance in their prediction to estimate uncertainty.   \n* Can the adaptive step size be implemented to generate images? If yes, what are the complications that may arise?\n\nI am willing to raise my score if the authors address all of my concerns.", "limitations": "The authors did not sufficiently detail all the limitations of their algorithm - For example, not emphasizing the additional computation time and memory involved in training the variance network, and the additional inference steps compared to Rectified Flow.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper shows an interesting idea - to use an estimation of the uncertainty in action prediction (the conditional variance) to dynamically adapt the step size of the flow model. If the model is certain in its action prediction, the step size will be large. Otherwise, the model performs many small prediction steps. \n* AdaFlow speeds-up the inference step of flow-based models, which is an important factor for many real-time applications. \n* The authors show the advantages of the proposed approach in various environments (4 tasks).\n* The paper is well-written and easy to follow and understand.", "weaknesses": "Major: \n* The main weakness of the paper is the lacking representation of the advantages/disadvantages of the main idea compared to Rectified Flow and Consistency models: The paper presents AdaFlow as the leading algorithm that reduces the inference time of diffusion policy, but it does not provide the full picture. As described in lines 199-206 (Discussion of AdaFlow and Rectified Flow), at inference time Rectified Flow allows for one-step action generation, which is either better than (for high variance actions) or equal (for low variance actions) to Adaflow. At training time, Rectified Flow involves a “reflow” stage after training the ODE. The reflow stage consists of generating pseudo-noise-data pairs and retraining the ODE. When comparing it to the training time of AdaFlow, AdaFlow also needs a two-stage training scheme for training the variance network (as described in lines 157-160), or a single, yet more involved training stage where the variance network is trained simultaneously with the ODE (which can also reduce performance). So both Rectified Flow and AdaFlow train two networks, the only difference is the generation of pseudo-noise-data pairs in Rectified Flow, which can be less time-consuming (or even negligible) compared to training the networks, depending on how many pairs are needed to get an accurate prediction.\\\n\\\nOverall, it seems that AdaFlow gains a reduction in training time (compared with Rectified Flow which generates pseudo-noise-data pairs) at the expense of longer inference time (due to additional steps for high variance actions). This means that it is not clear which algorithm is superior, and in general, it depends on the target task and dataset used.\\\n\\\nIn addition, line 204 mentions that the reflow stage leads to inferior generation quality due to straightened ODEs, but the reason for this claim is unclear and lacks intuition, since training a network for variance prediction as done in AdaFlow can also lead to inaccuracies and suboptimality in some cases (also Tables 2 and 4 show that Rectified Flow performs better or equal to AdaFlow in about half of the environments examined).\n\n\n\n* Comparison to the state of the art: The paper compares AdaFlow to the classical diffusion policy that uses DDIM and shows that AdaFlow achieves a higher success rate with less NFE. For example, Figure 3 shows that AdaFlow with NFE of 1.75 achieves comparable performance to a diffusion model with NFE of 100. To see the whole picture, it would be beneficial to compare AdaFlow to more advanced and faster solvers such as Dpm-solver and Dpm-solver++ (references 41 and 42 of the paper).\n\nMinor mistakes: \n* Figure 1 is not referenced in the paper and also the caption does not include an explanation of its subfigures. \n* Line 175 - larger instead of large\n* Table 5 - Maze 2 instead of Maze 1", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer df3a for the thoughtful questions and suggestions.\n\n**1. What is the advantage of AdaFlow over Rectified Flow with reflow?**  \n\n- **AdaFlow can learn action generation and variance estimation together, as shown in Appendix A.6**, where joint and separate training yield similar performance. We opted for separate training due to better efficiency and faster convergence, given the stationary objective of the variance network. Training the variance network is much quicker than training the original flow, adding minimal overhead to the 1-Rectified Flow. **In contrast, reflow and consistency distillation always require an expensive ODE distillation stage.**\n\n- Reflow is more complex and computationally expensive than AdaFlow. Reflow requires generating pseudo data by simulating the full trajectory (100 model inferences per pseudo sample) and needs a reflow dataset as large as the original. Then, it trains a new flow model with a reflow target, which can introduce errors. AdaFlow avoids this costly data generation. Training the variance estimation part separately (or jointly) involves only a few linear layers on the original rectified flow model. In practice, **on the maze task, reflow takes 7.5 times longer than AdaFlow in terms of wall clock time.**\n\n**2. Why does reflow lead to inferior generation quality with straightened ODE?** \n\n- Reflow estimates a **new mapping ODE** from noise to data, introducing errors. The first error comes from time-discretization during numerical simulation, and the second from training on pseudo data pairs. AdaFlow avoids these issues by using the original ODE.\n\n- In Tables 2 and 4, Reflow's rare outperformance of AdaFlow isn't due to variance prediction errors but a balance of efficiency and success rate. With more inference steps, AdaFlow outperforms Reflow, demonstrating a better ODE. The following tables compare Reflow and AdaFlow with different inference steps (NFE) on LIBERO tasks.\n\n\t|   | Task 1 | Task 2 | Task 3 | Task 4 | Task 5 | Task 6 | Average |\n\t| - | ------ | ------ | ------ | ------ | ------ | ------ | ------- |\n\t| Reflow  | 0.90 | 0.82  | 0.98   | 0.82   | 0.82   | 0.96   |  0.88   |\n\t| AdaFlow (NFE=1.27) | 0.98 | 0.80  | 0.98   | 0.82   | 0.90   | 0.96  |  0.91 |\n\t| AdaFlow (NFE=2.91) | 0.96 | 0.86  | 1.00   | 0.82   | 0.90   | 0.98  |  0.92 |  \n\n**3. Comparison to DPM solver and DPM ++ solver**\n\nThank you for suggesting a comparison with advanced solvers like DPM-solver and DPM-solver++. We conducted additional experiments, and as shown in the table below, AdaFlow achieves higher success rates even with fewer NFE across multiple LIBERO tasks.\n\n|   | NFE | Task 1 | Task 2 | Task 3 | Task 4 | Task 5 | Task 6 | Average |\n| - | ---   | ------ | ------ | ------ | ------ | ------ | ------ | ------- |\n| DPM     | 1    | 0.00 | 0.00  | 0.00   | 0.00   | 0.00   | 0.00  |  0.00 |\n| DPM     | 2    | 0.42 | 0.44  | 0.54   | 0.36   | 0.58   | 0.38  |  0.45 |\n| DPM     | 20   | 1.00 | 0.80  | 0.98   | 0.84   | 1.00   | 0.88  |  0.92 | \n| DPM++   | 1    | 0.00 | 0.00  | 0.00   | 0.00   | 0.00   | 0.00  |  0.00 |\n| DPM++   | 2    | 0.42 | 0.44  | 0.54   | 0.36   | 0.58   | 0.36  |  0.45 |\n| DPM++   | 20   | 1.00 | 0.80  | 0.98   | 0.84   | 1.00   | 0.92  |  0.92 | \n| AdaFlow | 1.27 | 0.98 | 0.80  | 0.98   | 0.82   | 0.90   | 0.96  |  0.91 |\n\n**4. Difference to RK45 solver**\n\nThe RK45 solver adjusts step size for accuracy by using the difference between its 4th and 5th order estimates to calculate an error estimate, which then adjusts the next step size.\n\t\nAdaFlow and RK45 are orthogonal methods for step size estimation:\n- RK45 increases numerical simulation accuracy.\n- AdaFlow leverages the straightness of the learned flow.\n\nThus, they can be used together. **The theoretical lowest NFE for RK45 is 6**, as it calculates 6 intermediate values per step, while AdaFlow needs close to one step.\n\n**5. Use other non-post-training methods to estimate the state variance**\n\nWe explored estimating state variance without training a variance estimation network. We trained only the rectified flow model and estimated state variance by sampling. For a given state, we sampled $N_z$ noises and $N_t$ timesteps, and then estimate the variance of each state following Equation 6. \n\n$$\\frac{1}{N_t} \\frac{1}{N_z} \\sum_t\\sum_{z_0} \\mathbb{E}[||y - z_0 - v(z_t, t; x)||^2]$$\n\nHere $y$ is derived using the Euler sampler, and $z_0$ is random Gaussian noise. This method estimates state variance accurately if both $N_z$ and $N_t$ is large enough. However, it involves sampling the full flow trajectory for many different noises. The figure below shows how many samples are needed to achieve a precise estimation of state variance: [Figure Link](https://anonymous.4open.science/r/adaflow_rebuttal-C0D6/variance_estimation.png). In general, an NFE > 10 is required to estimate the variance well. Compared to this method, AdaFlow is much faster for inference, only requiring one model inference. \n\n**6. Usage in generating images**\n\nGreat point! We've also tried using AdaFlow for image generation. Applying the learned variance network at inference does speed up the process. However, the speedup is limited because the learned flow is not straight due to the high dimensionality and variability in image generation.\n\nIn robotics tasks, many actions are deterministic given the states, unlike in image generation. For text-conditioned images, even detailed prompts can result in many possible images, making the generation trajectory non-deterministic. Thus, AdaFlow is better suited for robotics, where deterministic generations significantly reduce inference costs.\n\n**7. Typos**\nWe will correct the typos mentioned by the reviewer and will reflect these changes in the revised version.\n\nPlease let us know if you have any further questions or concerns. We hope that our rebuttal addresses all your points thoroughly, and we would greatly appreciate your reconsideration of the score. Thank you!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TotNf6am1z", "reviewer": "Reviewer_FRk8", "summary": "In this paper, the authors propose a new step schedule for sampling from a gradient flow that they claim accelerates the sampling process.    While they focus on the setting of Imitation Learning, in principle, their algorithm samples from some conditional distribution $\\pi_E$ of actions given observations on Euclidean space by first sampling $z_0$ a standard Gaussian and then evolving $z_t$ according to a gradient flow with velocity field $v_\\theta$ learned from data.  As the flow itself has to be discretized, there is a natural question as to how the discrete steps should be chosen.  Whereas many prior works rely on an Euler discretization with uniform step sizes, the authors propose using the conditional variance of the deviation of the action from the endpoint of a linearized flow as a measure of how large a step to take.  Under extremely strong assumptions, the authors demonstrate that their step size leads to control under squared Wasserstein distance between samples from their discretized flow and the desired sample.  They then empirically demonstrate the efficacy of their techniques in a variety of environments.", "questions": "See weaknesses.", "limitations": "The authors do not address the limitations of their approach, despite the fact that estimating the conditional variance is a potentially challenging problem and their Gaussianity assumption in their suggested approach is potentially limiting.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper is concerned about the inference speed of learned imitator policies, which is an important problem that is sometimes ignored in academic work on IL, especially with the recent focus on diffusion models.  The authors present a diverse set of examples that is somewhat suggestive of the efficacy of their techniques.", "weaknesses": "The \"theoretical analysis\" (Line 59) presented is imprecise and holds only under extremely strong and unrealistic conditions.  As an example of the imprecision, the conditioning in the definition of the conditional variance in line 147 is a bit confusing to me.  Is the conditioning on $x_t$ intended to be conditional on the path up to time $t$?  If not, then it is not obvious to me that we could expect that $\\sigma^2$ can be zero in any setting where the gradient flow converges, due to the variance coming from $x_0$.  If so, this should be formally defined.  Furthermore, many of the claimed propositions, such as 3.1 and 3.2 are immediate.  Proposition 3.3 has a non-rigorous proof that refers to a vague property of a rectified flow without stating precisely what this property is or citing the result.  While in principle, any estimator of the conditional variance would be acceptable, the authors only present one that assumes the conditional distributions of the actions are Gaussian with no justification and no discussion of this assumption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We believe the reviewer might have some misunderstandings, and we will ensure to further clarify our statements in the updated version of the paper.\n\n- **In Line 147**, $x_t$ is a random variable that follows the marginalized probability at time $t$. Thus, $x_t$ is not conditional on the path up to time $t$ (similar to the original Rectified Flow paper). The conditional variance $\\text{Var}(a - x_0 | x_t = x, s)$ can indeed be zero, even though $x_0$ is random. This is because the variance is conditional on $x_t = x$. For instance, in a 1-dimensional example where $a = 1$, $x_0$ follows a Gaussian distribution $N(0, 1)$, and we define $x_t = (1 - t) x_0 + t$, as in Rectified Flow. In this case, given $x_t = x$, we can deterministically solve $(a - x_0)$ because $x_0 = \\frac{x_t - t}{1 - t}$ and $a$ is a constant 1. In fact, in this simple 1D example, the flow is exactly straight because $a$ is deterministic, which is what Proposition 3.2 is stating.\n\n- **Regarding Proposition 3.3**, the statement is accurate and correct. However, we will ensure that the statement is made more rigorously with a clear presentation of the assumptions and properties from earlier works.\n\n- The property mentioned is that $z_t$ and $x_t$ share the same law (i.e., for any $t$, their distributions are the same), which is formally stated as Theorem 3.3 in the Rectified Flow paper.\n\n- We respectfully note that we do not assume the conditional distributions of the actions are Gaussian. The action can follow any distribution, similar to the image distribution in image diffusion applications. Only the initial noise $x_0$ follows a Gaussian distribution. We would appreciate it if the reviewer could clarify what specific assumption or detail they are referring to regarding the Gaussian distribution.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ea7EDRVjTF", "reviewer": "Reviewer_AbSH", "summary": "This paper introduces AdaFlow, a Flow-based policy for imitation learning. The key innovation is a variance-adaptive solver that employs smaller step sizes when complexity is high and reduces to one-step generation when straightforward. The authors provide theoretical analysis, offering an error bound based on a threshold hyperparameter η. Experiments on maze navigation and robot manipulation tasks demonstrate AdaFlow's effectiveness, generating higher success rates and more diverse trajectories compared to popular Diffusion Policies, while requiring significantly fewer NFEs", "questions": "- In Table 2 and Table 8, you mention a \"diversity score,\" but I cannot find it in the tables. Could you clarify this discrepancy and explain how you quantify the diversity of generated trajectories?\n- How many seeds did you use to train the policies, and how many runs did you use to evaluate each policy? This information seems to be missing from the text.\n- How does AdaFlow compare with classical adaptive ODE solvers?\n- While AdaFlow doesn't need to perform reflow as in Rectified Flow, it requires training an extra variance estimation network. Could you provide a comparison of the computational costs for training these two methods?\n- What is the robustness of AdaFlow to potential errors in the variance estimation network?", "limitations": "No, the author doesn't have a discussion about the limitation. As I mentioned in the weakness section, I think the method may have some problems to scale to higher-dimensional action space where learning a ODE and the variance estimation won't be so easy.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper proposes a simple yet efficient adaptive scheme to dramatically reduce the computational cost for inference in generative model policies.\n- The method is well-grounded in theoretical analysis.\n- The paper is generally well-written and clearly structured.", "weaknesses": "- The environments tested in this work are standard but relatively low-dimensional. It's unclear how well the proposed method scales to high-dimensional tasks like humanoid control.\n- The experimental section could benefit from more detailed information about the setup and evaluation metrics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback, which has helped us improve the clarity and comprehensiveness of our work.\n\n**Possibility of extending to higher-dimensional robotics tasks**\n\nWe agree that applying AdaFlow to more complex and high-dimensional tasks like humanoid control would be interesting. Due to the time constraints of the rebuttal period, we couldn't include such experiments. However, on the LIBERO dataset, our input is video data plus the gripper configuration of the robot hand, which is one of the most high-dimensional tasks in robot manipulation. In future work, we plan to explore using AdaFlow on humanoid tasks. We are confident that training an ODE on higher-dimensional spaces is feasible, as ODE training has been successful in image generation tasks (as in Stable Diffusion 3)\n\n**Detailed experiment setup**\n\nWe will update the paper to include more detailed information about the experimental setup and evaluation metrics in Appendix A.9 Additional Experimental Details. Additionally, we will provide the codebase to assist with reproducing the results. \n\n**Correction of typos**\n\nWe have corrected the typos in the paper to ensure consistency and will reflect these changes in the revised version.\n\n**Seeds and runs**\n\nDue to the limited time during the rebuttal period, we couldn't complete the experiments with multiple seeds and runs. However, the partial results we obtain indicate that the performance of AdaFlow, diffusion policy, and other baselines are consistent with what we presented in the paper. We will include the results with multiple seeds in the revision to provide a comprehensive evaluation.\n\n**Comparison with classical adaptive ODE solvers**\n\nWe conducted experiments to compare AdaFlow with recent advanced and faster solvers like DPM-solver and DPM-solver++. As shown in the table below, we compared AdaFlow with DPM and DPM++ across multiple tasks on LIBERO. The results demonstrate that AdaFlow achieves higher success rates even with fewer NFE:\n\n|   | NFE | Task 1 | Task 2 | Task 3 | Task 4 | Task 5 | Task 6 | Average |\n| - | ---   | ------ | ------ | ------ | ------ | ------ | ------ | ------- |\n| DPM     | 1    | 0.00 | 0.00  | 0.00   | 0.00   | 0.00   | 0.00  |  0.00 |\n| DPM     | 2    | 0.42 | 0.44  | 0.54   | 0.36   | 0.58   | 0.38  |  0.45 |\n| DPM     | 20   | 1.00 | 0.80  | 0.98   | 0.84   | 1.00   | 0.88  |  0.92 | \n| DPM++   | 1    | 0.00 | 0.00  | 0.00   | 0.00   | 0.00   | 0.00  |  0.00 |\n| DPM++   | 2    | 0.42 | 0.44  | 0.54   | 0.36   | 0.58   | 0.36  |  0.45 |\n| DPM++   | 20   | 1.00 | 0.80  | 0.98   | 0.84   | 1.00   | 0.92  |  0.92 | \n| AdaFlow | 1.27 | 0.98 | 0.80  | 0.98   | 0.82   | 0.90   | 0.96  |  0.91 |\n\n\n**Robustness to potential errors in the variance estimation network**\n\nWe acknowledge that the variance estimation network may have errors. As shown in [this figure](https://anonymous.4open.science/r/adaflow_rebuttal-C0D6/variance_error.png), the predicted variance does not perfectly match the ground truth state variance. This estimation error can potentially influence the generation results. \n\nHowever, note that in Algorithm 1, $\\epsilon_t$ is clipped between $\\epsilon_\\text{min}$ and $(1 - t)$, so in theory, AdaFlow’s worst case performance (with arbitrarily bad variance estimation) will be between 1-step RF and $1 / \\epsilon_\\text{min}$-step RF.\n\nEmpirically, we compare AdaFlow against RF with 20 inference steps on the LIBERO task suites, the results are shown in the following:\n\n|   | NFE | Task 1 | Task 2 | Task 3 | Task 4 | Task 5 | Task 6 | Average |\n| - | --- | ------ | ------ | ------ | ------ | ------ | ------ | ------- |\n| Rectified Flow | 20   | 0.96 | 0.86  | 1.00   | 0.82   | 0.86   | 0.98  |  0.92 |\n| AdaFlow        | 1.27 | 0.98 | 0.80  | 0.98   | 0.82   | 0.90   | 0.96  |  0.91 |\n\nAs shown, AdaFlow's performance is comparable to Rectified Flow with NFE=20. This demonstrates that the introduction of the variance estimation process does not significantly impact performance. AdaFlow remains effective and efficient.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "onT41gnQNn", "reviewer": "Reviewer_RJmb", "summary": "This paper focuses on accelerating flow-based imitation learning algorithms by extending the existing work of Rectified Flow to the field of imitation learning, proposing AdaFlow. The paper leverages probability flows to train a state-conditioned action generator and introduces a variance-adaptive ODE solver to adjust the step size, accelerating inference while ensuring diversity. Experimental results in four environments, including navigation and manipulation, demonstrate the efficiency of AdaFlow.", "questions": "- Does the step size depend not only on the variance (number of potential behaviors corresponding to the same state) but also on the complexity of expert behaviors?\n- In Tables 2, 3, and 4, Reflow achieves comparable (sometimes better) success rates to AdaFlow with NFE=1, which raises doubts about the advantages of AdaFlow. Additionally, can Reflow+distillation achieve even better results as described in [1]?\n- The top and bottom in Fig. 6 do not match the description in line 269. Furthermore, as shown in Fig. 6, Diffusion Policy achieves SR=1 at NFE=3~5, contradicting the claim of \"requiring 10x less inference times to generate actions.\"\n- The experiments only use success rate as an evaluation metric; however, commonly used metrics in clustering algorithms, such as NMI and ENT, can be considered as measures of behavioral richness.\n- What does \"Hz\" mean in Table 1?\n- In formula (7), the log likelihood should be $\\log(\\sigma^2)$ instead of $\\log(\\sigma)$.\n- Please carefully check and correct the formatting of references, e.g., [36] lacks author information, and references like [15], [32], [34] lack dates.\n\n[1] Liu X, Gong C, Liu Q. Flow straight and fast: Learning to generate and transfer data with rectified flow[J]. arXiv preprint arXiv:2209.03003, 2022.", "limitations": "The paper lacks a discussion on the limitations of the method.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The idea of this paper is well-motivated and the investigated problem is practical.\n- The theme of this paper in speeding up flow-based generative models is timely and interesting.\n- Theoretical analysis is complete.\n- Empirical results show improvement over baseline and other compared methods. The visualized results are good to demonstrate the efficiency.", "weaknesses": "- There are some minor typos in the paper.\n- The implementation codes are not provided.\n- The experimental data lack statistical significance.\n- The limitations of the method are not discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback, which has helped us improve the clarity and comprehensiveness of our work.\n\n**Complexity of expert behavior and variance**\n\nThe complexity of expert behavior and the variance are intertwined. Since we perform offline learning with demonstration data, the model learns a high variance when the expert’s behavior becomes diverse at a given state, and vice versa. So even if the expert behavior is very complicated (like peg in the hole), but the action itself is relatively deterministic (there is a consistent way of solving the underlying task), then AdaFlow will still choose a large step size, because the model is more certain about which action to take, so it performs more like in the BC case. In contrast, even if it is a simple task like lifting an object, if there exists many different ways of lifting, then AdaFlow will still choose a small step size, until the lifting behavior is determined (the conditional variance drops accordingly).\n\n**Comparison with Reflow**\n\nYes, Reflow can achieve comparable success rates to AdaFlow. However, Reflow requires an additional complex and time-consuming stage after training the original flow. AdaFlow is much more efficient than Reflow. This is because: \n\n- AdaFlow doesn’t have to be trained in two separate stages, **it can also learn action generation and variance estimation together**. As shown in the following table (originally in Appendix A.6), training the flow and the variance networks separately or jointly results in similar performance. In practice, we chose to train them separately because separate training in fact results in even better training efficiency (faster convergence), this is because the objective of the variance network is stationary. In addition, training the variance network is way faster than training the original flow, hence Adaflow introduces almost no overhead to the 1-Rectified Flow (without reflow). **Compared to AdaFlow, reflow and consistency distillation always require a stage to distill the ODE，which itself alone is more expensive than the original training phase.**\n\n\t|  | Maze 1 | Maze 2 | Maze 3 | Maze 4 |\n\t| - | ------ | ------ | ------ | ------ |\n\t| AdaFlow (Separate) | 0.98 | 1.00 | 0.96 | 0.86 |\n\t| AdaFlow (Joint)    | 1.00 | 1.00 | 0.96 | 0.88 |\n\n- Reflow is more complicated and computationally expensive than AdaFlow. Reflow operates in two steps: first, it generates pseudo data, requiring simulating the full trajectory, i.e., 100 model inferences per pseudo sample, and the reflow dataset usually needs to be as large as the original dataset. Second, it trains the new flow model with a reflow target, which might introduce errors. AdaFlow doesn’t need this expensive data generation. If we train the variance estimation part separately (it can also be trained jointly), it only involves training a few linear layers on top of the original rectified flow model. In practice, we compared the total training time for the reflow stage and AdaFlow’s variance estimation training stage. On the maze task, we observed that reflow takes **7.5x longer than AdaFlow in terms of wall clock time.**\n\n- It is possible to apply a distillation stage after reflow. In image and video generation tasks, distillation is often necessary because the ODE might not be straight even after reflow. However, in our robotics tasks, reflow has already learned a very straight trajectory. We observed no further improvements when using more than one inference step for reflow. Therefore, the potential improvement from distillation is very minimal. Additionally, distillation requires another training stage, which adds complexity and inconvenience.\n\n**Measure behavior richness**\n\nWe appreciate the suggestion to use NMI or ENT as a measure of behavioral richness. However, it's not feasible to apply them directly to the generated trajectories because each trajectory has a different length. Clustering methods like NMI typically require fixed-length data, which isn't the case with our generated trajectories. We are open to exploring alternative metrics that could better capture the diversity and richness of behaviors in future work. If you have any suggestions or additional metrics that you believe could be useful, we would be happy to consider them.\n\n**Typos**\n\nThere’s a typo in Table 1, which we will correct in the revised version. We will correct other typos mentioned by the reviewer in the paper and will reflect these changes in the revised version.\n\nPlease let us know if you have any further questions or concerns. We hope that our rebuttal addresses all your points thoroughly, and we would greatly appreciate your reconsideration of the score. Thank you!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The work develops a flow simulation method in the context of imitation learning using a flow-based generative model. The method is based on the connection between variance and flow straightness, where the variance is estimated by an additional model which can be trained together with the flow model, which is then used to adjust discretization step size when interpreted as an indicator of flow curvature. All reviewers acknowledged the novel contribution and the significance to the field. Reviewers also raised a few questions and concerns, most of which have been reasonably addressed, while there also remain some unclear points, e.g., the applicability to high-dimensional problems (like image generation). Nevertheless, the paper has made a reasonable contribution, and is recommended for acceptance.
Decision: Accept (poster)