Title: AdaFlow: Imitation Learning with Variance-Adaptive Flow-Based Policies
Venue: NeurIPS 2024
Accepted: yes
Keywords: Imitation Learning; Generative Model-Based Policy
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

Abstract

Diffusion-based imitation learning improves Behavioral Cloning (BC) on multimodal 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 flowbased 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.

Noise

Step 50

Step 100

AdaFlow

Noise

Action

Diffusion

Policy

Noise

Step 50

Step 100

Noise

Low-Variance State

High-Variance State

Step 5

Step 10

1

Low

High

1 Introduction

Imitation Learning (IL) is a widely adopted method in robot learning [1, 2]. In IL, an agent is given a demonstration dataset from a human expert finishing a certain task, and the goal is for it to complete the task by learning from this dataset. IL is notably effective for learning complex, non-declarative motions, yielding remarkable successes in training real robots [3-6].

𝑡

𝑡

𝑡

𝑡

𝑡

Data

Prediction

𝑡

𝑡

𝑡

𝑡

𝑡

𝑥

𝑥

𝑥

𝑦

𝑥 = -5

𝑥 = -2.5

𝑥 = 0

𝑥 = 2.5

𝑥 = 5

Ground Truth

Diffusion Policy

AdaFlow (Euler)

AdaFlow

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, 11, 6, 12]. To improve it, various frameworks have been proposed. For instance, Implicit BC [12] learns an energy-based model for each state and searches the actions that minimize the energy with optimization algorithms. Diffuser [13, 14] and Diffusion Policy [6] adopts diffusion models [15, 16] to generate diverse actions, which has become the default method for training on large-scale robotics data [17-20].

The computational cost of the learned policies at the execution stage is important for an IL framework in a real-world deployment. Unfortunately, none of the previous frameworks enjoys both efficient inference and diversity. Although energy-based models and diffusion models can generate multimodal 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, we propose a new IL framework, named AdaFlow, that learns a dynamic generative policy that can autonomously adjust its 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 are powerful generative models that precisely capture the complicated distributions, but similar to energy-based models and diffusion models, they still require multiple recursive iterations to 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. In contrast, both of these methods require an additional distillation or 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 as naive BC.

Empirical results across multiple benchmarks demonstrate that AdaFlow achieve consistently good performance across success rate with high execution efficiency. Specifically, our contributions are:

Weproposed AdaFlow as a generative model-based policy for decision making tasks, capable of generating actions almost as quickly as a 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 [10]. AdaFlow consistently outperforms existing state-of-the-art models, despite requiring 10x less inference times to generate actions.

Weoffer a theoretical analysis of the overall error in action generation by AdaFlow, providing a bound that ensures precision and reliability.

2 Related Work

Diffusion/Flow-based Generative Models and Adaptive Inference Diffusion models [28, 16, 15, 29] succeed in various applications, e.g., image/video generation [30-33], audio generation [34], point cloud generation [35-38], 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-44] by transforming the diffusion process into marginal-preserving 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 Runge-Kutta [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 decision-making tasks in real-world applications. Recently, new generative models [21, 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 to their preference of straight trajectories. Among them, Rectified flow achieves one step generation with reflow, a 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 [48-51]. In a pioneering work, Janner et al. [13] proposed 'Diffuser', a planning algorithm with diffusion models for offline reinforcement learning. This framework is 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 to model policies as conditional diffusion models. The application of DDPM [16] and DDIM [43] 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 a 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 propose 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 a state before allocating

computation. If the state has a 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 is made possible by leveraging a combination of elements: 1) a special property of the flow 2) a 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 = { ( s ( i ) , a ( i ) ) } n i =1 , our goal is to learn a policy π θ that can generate trajectories following the target distribution π E . π θ can be induced from a state-conditioned flowbased model,

Here, s is the state and the velocity field is parameterized by a neural network θ 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 a state-conditioned least-squares objective,

where x t is the linear interpolation between x 0 and x 1 = a :

We should differentiate z t which is the ODE trajectory in (1) from the linear interpolation x t . They do not overlap unless all trajectories of ODE (1) are straight. See Liu et al. [21] for more discussion.

With infinite data sampled from π E , unlimited model capacity and perfect optimization, it is guaranteed that the policy π θ generated from the learned flow matches the expert policy π E [21].

3.2 The Variance-Adaptive Nature of Flow

Typically, to sample from the distribution π θ at state s , we start with a random sample z 0 from the Gaussian distribution and simulate the ODE (Eq. (1)) with multi-step ODE solvers to get the action. For example, we can exploit N -step Euler discretization,

After running the solver, z 1 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 a 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 is chosen, the subsequent actions to take will be almost deterministic. This renders using a uniform Euler solver with the same number of inference steps N across all the states a sub-optimal solution. Rather, it is preferred that the agent can vary its decision-making process as the state of the agent changes. The challenge is how to quantitatively estimate the complexity of a 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 a state with the training loss of the flow-based policy: if the distribution of actions is deterministic at a state s (i.e., a Dirac distribution), the trajectory of rectified flow ODE is a straight line, i.e., a single Euler step yields an exact estimation of z 1 .

Proposition 3.1. Let v ∗ be the optimum of Eq. (2) . If var π E ( a | s ) = 0 where a ∼ π E ( · | s ) , then the learned ODE conditioned on s is

whose trajectories are straight lines pointing to z 1 and hence can be calculated exactly with one step of Euler step:

Algorithm 1 AdaFlow: Execution

1: Input: Current state s , minimal step size ϵ min, error threshold η , pre-trained networks v θ and σ ϕ .

2: Initialize z 0 ∼ N (0 , I ) , t = 0 .

3: while t < 1 do

4: Compute step size

5: Update t ← t + ϵ t , z t ← z t + ϵ t v θ ( z t , t | s ) .

6: end while

7: Execute action a = z 1 .

Note that the straight trajectories of (5) satisfies z t = t a +(1 -t ) z 0 , which makes it coincides with the linear interpolation x t . As show in [21], this happens only when all the linear trajectories do not intersect on time t ∈ [0 , 1) .

More generally, we can expect that the straightness of the ODE trajectories depends on how deterministic the expert policy π E is. Moreover, the straightness can be quantified by a conditional variance metric defined as follows:

Proposition 3.2. Under the same condition as Proposition 3.1, we have σ 2 ( z t , t | s ) = 0 from (5) .

The proof of the above propositions is in Appendix A.1. To summarize, the variance of the stateconditioned loss function at ( z t , 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 a huge amount of computation. In Section 3.3, we will show the variance can be used to bound the discretization error, thereby enabling the design of an adaptive ODE solver.

Variance Estimation Network In practice, the conditional variance σ 2 ( x, t | s ) can be empirically approximated by a neural network σ 2 ϕ ( x, t | s ) with parameter ϕ . Once the neural velocity v θ is learned, we can estimate σ ϕ by minimizing the following Gaussian negative log-likelihood loss:

We adopt a two-stage training strategy by first training the velocity network v θ then the variance estimation network σ ϕ . In practice, the second stage just involves fine-tuning a 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-Adaptive Flow-Based Policy

Because the variance indicates the straightness of the ODE trajectory, it allows us to develop an adaptive approach to set the step size to yield better estimation with lower error during inference.

To derive our method, let us consider to advance the ODE with step size ϵ t at z t :

The problem is how to set the step size ϵ t properly. If ϵ t is too large, the discretized solution will significantly diverge from the continuous solution; if ϵ t is too small, it will take excessively many steps to compute.

We propose an adaptive ODE solver based on the principle of matching the discretized marginal distribution p t of z t from (8), and the ideal marginal distribution p ∗ t when following the exact ODE

(1). This is made possible with a key insight below showing that the discretization error can be bounded by the conditional variance σ 2 ( z t , t | s ) .

Proposition 3.3. Let p ∗ t be the marginal distribution of the exact ODE d z t = v ∗ ( z t , t | s )d t . Assume z t ∼ p t = p ∗ t and p t + ϵ t the distribution of z t + ϵ t following (8) . Then we have

where W 2 denotes the 2-Wasserstein distance.

We provide the proof in Appendix A.2. Hence, given a threshold η , to ensure that an error of W 2 ( p ∗ t + ϵ t , p t + ϵ t ) 2 ≤ η 2 , we can bound the step size by ϵ t ≤ η/σ ( z t , t | s ) . Because ϵ t at time t should not be large than 1 -t , we suggest the following rule for setting the step size ϵ t at z t at time t :

where we impose an additional lower bound ϵ min to avoid ϵ t to be unnecessarily small. Besides, the proposed adaptive strategy guarantees to instantly arrive at the terminal point when σ 2 ( z t , t | s ) = 0 as ϵ t = 1 -t . Moreover, it aligns with Section. 3.2 since for states with deterministic actions, it sets ϵ 0 = 1 to generate the action in one step. We incorporate the above insights to the execution in Algorithm 1.

Global Error Analysis Proposition 3.3 provides the local error at each Euler step. In the following, we provide an analysis of the overall error for generating z 1 when we simulate ODE while following the adaptive rule (9). To simplify the notation, we drop the dependency on the state s , and write v ∗ t ( · ) = v ∗ ( · , t | s ) .

Assumption 3.4. Assume ∥ v ∗ t ∥ Lip ≤ L for t ∈ [0 , 1] , and the solutions of d z t = v t ( z t )d t has bounded second curvature ∥ ¨ z t ∥ ≤ M for t ∈ [0 , 1] .

This is a standard assumption in numerical analysis, under which Euler's method with a constant step size of ϵ min admits a global error of order O ( ϵ min ) .

Proposition 3.5. Under Assumption 3.4, assume we follow Euler step (8) with step size ϵ t in (9) , starting from z 0 = x 0 ∼ p ∗ 0 . Let p t be the distribution of z t we obtained in this way, and p ∗ t that of x t in (3) . Note that p ∗ 1 is the true data distribution. Set η = M η ϵ 2 min / 2 for some M η > 0 , and ϵ min = 1 /N max .

Let N ada be the number of steps we arrive at z 1 following the adaptive schedule. We have

where C is a constant depending on M , M η and L .

The idea is that the error is proportional to N ada N max , suggesting that the algorithm claims an improved error bound in the good case when it takes a smaller number of steps than the standard Euler method with constant step size ϵ min . 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, our method utilizes only the original ODE, eliminating the need for an 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) a 2D navigation problem; and two robot

manipulation task suites on 3) RoboMimic [14] following past works [6] and 4) LIBERO [3], provide diverse and realistic scenarios for evaluation.

Our results show that AdaFlow improves the success rate of completing both navigation and manipulation tasks, outperforming state-of-the-art methods such as BC and its variants, as well as Diffusion Policy, across a 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.

4.1 Regression

We start with a 1D regression task designed to demonstrate the adaptivity nature of AdaFlow. The goal is to learn a mapping from x to y where

Note that y | x is deterministic when x ≤ 0 and stochastic otherwise. The training and testing data are uniformly sampled from the ground-truth function with x ∈ [ -5 , 5] . Details about the setup and the hyperparameters are provided in Appendix.

AdaFlow can achieve 1-step generation for deterministic states. Figure 2 (top-right) shows the generation trajectories of Diffusion Policy and AdaFlow with 5 step. Notably, when x ≤ 0 , AdaFlow generates straight trajectories and is therefore able to predict y with a single step, aligning our analysis in Proposition 3.1 and 3.2. In contrast, Diffusion Policy generates curved trajectories when step = 5, and hence cannot predict y accurately with a single step. The bottom of Figure 2 shows the estimated variance by AdaFlow across x ∈ [ -5 , 5] , which accurately aligns with the expected variance. In addition, as x increases, AdaFlow adaptively increases the required number of simulation steps.

4.2 Navigating a 2D Maze

We create two sets of maze navigation tasks to validate AdaFlow's performance of modeling multimodal behavior. In particular, we create two single-task environments where the agent starts and ends at a 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.

AdaFlow

(NFE=1.12)

BC

(NFE=1)

π

(NFE=1)

π

End

Start

Put the black bowl at

the front on the plate

Put the middle black

bowl on the plate

Put the middle black bowl

on top of the cabinet

Stack the front black on the

black bowl in the middle

End

Start

Open the top drawer

of the cabinet

Put the black bowl at

the back on the plate

Task 1

Task 2

Task 3

Task 4

Task 5

Task 6

AdaFlow achieves high diversity and success with low NFE; Diffusion Policy and BC lag in comparison. We compare AdaFlow against baseline methods in Table 2. We additionally visualize the rollout trajectories from each learned policy in Figure 3 as a 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 can 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 to 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 on practical robotics tasks, we compare AdaFlow against baselines on a Push-T task [6], the RoboMimic [10] benchmark (Lift, Can, Square, Transport, ToolHang) and the LIBERO [27] 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, we add three additional baseline methods: 1) LSTM-GMM, BC with the LSTM model and a Gaussian mixture head, 2) IBC, the implicit behavioral cloning [12], an energy-based model for generative decision-making, and 3) BET [11]. For the LIBERO tasks, we pick a subset of six Kitchen tasks and follow the setup described in the LIBERO paper (Check Figure 4 for the description of the six tasks).

AdaFlow consistently outperforms competitors in 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 or higher success rates across different challenging manipulation tasks, compared against all baselines, with only an 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 1.27. We additionally visualize the variance predicted by AdaFlow in Figure 5. 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

1

0

DS

SR

SR

SR

NFE

DS

BC

Diffusion Policy

Epoch

NFE (

𝜂

)

SR

SR

DS

SR

Epoch

SR

SR

NFE

NFE

DS

BC

Diffusion Policy

AdaFlow (Ours)

1.00

0.75

0.50

0.25

0.00

1

5

8

20

1.00

F

0.75

0.25

0.00

0

100

150

200

Higher Training and Inference Efficiency. Figure 6 (top) examines changes in success rate relative to the NFE. AdaFlow maintains a 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 over epochs. It shows that AdaFlow has a better area-under-curve than Diffusion Policy, indicating faster learning. As expected, due to its simplicity, Behavioral Cloning (BC) achieves the best learning efficiency.

Robustness to η . In Figure 6, the NFEs in AdaFlow are calculated at various η values. It shows that AdaFlow is robust to changes in η .

On the Importance of Variance Estimation. In Table 5, we provide the performance of AdaFlow with and without the variance estimation network on the four mazes from Section 4.2. From the results, it is clear that the variance estimation network not only makes inference faster, but can also lead to better performance.

AdaFlow (Ours)

Epoch

NFE

5 Conclusion

We present AdaFlow, a 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 training/execution efficiency. This work lays a robust foundation for future research on adaptive imitation learning methods in real-world scenarios.

References

Schaal, S. Is imitation learning the route to humanoid robots? Trends in cognitive sciences , 3(6):233-242, 1999.

Osa, T., J. Pajarinen, G. Neumann, et al. An algorithmic perspective on imitation learning. Foundations and Trends® in Robotics , 7(1-2):1-179, 2018.

Liu, B., X. Xiao, P. Stone. A lifelong learning approach to mobile robot navigation. IEEE Robotics and Automation Letters , 6(2):1090-1096, 2021.

Zhu, Y., A. Joshi, P. Stone, et al. Viola: Imitation learning for vision-based manipulation with object proposal priors. 6th Annual Conference on Robot Learning (CoRL) , 2022.

Brohan, A., N. Brown, J. Carbajal, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06817 , 2022.

Chi, C., S. Feng, Y. Du, et al. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137 , 2023.

Pomerleau, D. A. Alvinn: An autonomous land vehicle in a neural network. Advances in neural information processing systems , 1, 1988.

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 statistics , pages 627-635. JMLR Workshop and Conference Proceedings, 2011.

Torabi, F., G. Warnell, P. Stone. Behavioral cloning from observation. In Proceedings of the Twenty-Seventh International Joint Conference on Artificial Intelligence . International Joint Conferences on Artificial Intelligence Organization, 2018.

Mandlekar, A., D. Xu, J. Wong, et al. What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108.03298 , 2021.

Shafiullah, N. M., Z. Cui, A. A. Altanzaya, et al. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing systems , 35:22955-22968, 2022.

Florence, P., C. Lynch, A. Zeng, et al. Implicit behavioral cloning. In Conference on Robot Learning , pages 158-168. PMLR, 2022.

Janner, M., Y. Du, J. Tenenbaum, et al. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning , pages 9902-9915. PMLR, 2022.

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.

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.

Ho, J., A. Jain, P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Sridhar, A., D. Shah, C. Glossop, et al. Nomad: Goal masked diffusion policies for navigation and exploration. arXiv preprint arXiv:2310.07896 , 2023.

Team, O. M., D. Ghosh, H. Walke, et al. Octo: An open-source generalist robot policy, 2023.

Shah, D., A. Sridhar, N. Dashora, et al. Vint: A foundation model for visual navigation. arXiv preprint arXiv:2306.14846 , 2023.

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.

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.

Lipman, Y., R. T. Chen, H. Ben-Hamu, et al. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations . 2022.

Albergo, M. S., N. M. Boffi, E. Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 , 2023.

Heitz, E., L. Belcour, T. Chambon. Iterative α -(de)blending: A minimalist deterministic diffusion model. In ACM SIGGRAPH 2023 Conference Proceedings , SIGGRAPH '23. Association for Computing Machinery, New York, NY, USA, 2023.

Liu, Q. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 , 2022.

Song, Y., P. Dhariwal, M. Chen, et al. Consistency models. 2023.

Liu, B., Y. Zhu, C. Gao, et al. Libero: Benchmarking knowledge transfer for lifelong robot learning. arXiv preprint arXiv:2306.03310 , 2023.

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.

Song, Y., S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Ho, J., T. Salimans, A. A. Gritsenko, et al. Video diffusion models. In Advances in Neural Information Processing Systems .

Zhang, S., X. Yang, Y. Feng, et al. Hive: Harnessing human feedback for instructional visual editing. arXiv preprint arXiv:2303.09618 , 2023.

Wu, L., C. Gong, X. Liu, et al. Diffusion-based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems , 35:36533-36545, 2022.

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-36494, 2022.

Kong, Z., W. Ping, J. Huang, et al. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations . 2020.

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.

Liu, X., L. Wu, M. Ye, et al. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations . 2023.

Luo, S., W. Hu. Score-based point cloud denoising. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 4583-4592. 2021.

Wu, L., D. Wang, C. Gong, et al. Fast point cloud generation with straight flows. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 9445-9454. 2023.

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.

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.

Liu, L., Y. Ren, Z. Lin, et al. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations . 2021.

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.

Song, J., C. Meng, S. Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations . 2020.

Bao, F., C. Li, J. Zhu, et al. Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. In International Conference on Learning Representations . 2021.

Press, W. H., S. A. Teukolsky. Adaptive stepsize runge-kutta integration. Computers in Physics , 6(2):188-191, 1992.

Albergo, M. S., E. Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations . 2022.

Salimans, T., J. Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations . 2022.

Kapelyukh, I., V. Vosylius, E. Johns. Dall-e-bot: Introducing web-scale diffusion models to robotics. IEEE Robotics and Automation Letters , 2023.

Yang, S., O. Nachum, Y. Du, et al. Foundation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.04129 , 2023.

Pearce, T., T. Rashid, A. Kanervisto, et al. Imitating human behaviour with diffusion models. In Deep Reinforcement Learning Workshop NeurIPS 2022 . 2022.

Chang, W.-D., J. C. G. Higuera, S. Fujimoto, et al. Il-flow: Imitation learning from observation using normalizing flows. arXiv preprint arXiv:2205.09251 , 2022.

Wang, Z., H. Zheng, P. He, et al. Diffusion-gan: Training gans with diffusion. In The Eleventh International Conference on Learning Representations . 2022.

Freund, G. J., E. Sarafian, S. Kraus. A coupled flow approach to imitation learning. In International Conference on Machine Learning , pages 10357-10372. PMLR, 2023.

Fu, J., A. Kumar, O. Nachum, et al. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

A Appendix

A.1 Proof of Proposition 3.1 and Proposition 3.2.

PROOF 1. var π E ( a | s ) = 0 means that the action a = a equals a deterministic value a given s . With x t = t a +(1 -t ) x 0 , note that

Therefore, a -x 0 is deterministically decided by x t and s . This yields

Differentiating it also yields

Therefore, we have

Solving ODE this yields

We also have σ 2 ( x, t | s ) again because a -x 0 is deterministic given x t and s :

A.2 Proof of Proposition 3.3

PROOF 2. Following the property of rectified flow, the distribution of x 1 = t a +(1 -t ) x 0 coincides with p t for all t ∈ [0 , 1] . Hence, we can assume that z t = x t ∼ p ∗ t . In this case, we have z t + ϵ t = x t + ϵ t v ∗ ( z t , t | s ) and x t + ϵ t = x t + ϵ t ( a -x 0 ) . We have

A.3 Proof of Proposition 3.5

PROOF 3. Assume the adaptive algorithm visits the time grid of 0 = t 0 , t 1 , . . . , t N = 1 .

Define z t i t be the result when we implement the adaptive discretization algorithm upto t i and then switch to follow the exact ODE afterward, that is, we have d z t i t = v t ( z t i t )d t for t ≥ t i . In this way, we have z 1 t = z t , and z 0 t = z ∗ t , where z ∗ t is the trajectory of the exact ODE d z ∗ t = v ∗ t ( z ∗ t )d t .

From Lemma A.1, we have

Let p t i t be the distribution of z t i t . Then we have p 1 t = p t and p 0 t = p ∗ t . Then

where C = 1 2 max( M,M η ) exp( L (1 -t i )) . Here we use the bound in the proof of Proposition 3.1 and Lemma A.1. Hence,

where C = exp( L ) max( M,M η ) .

Lemma A.1. Let ∥ v t ∥ Lip ≤ L for t ∈ [0 , 1] . Assume x t and y t solve d x t = v t ( x t )d t and d y t = v t ( y t )d t starting from x 0 , y 0 , respectively. We have

where we used ∥ v t ( x t ) -v t ( y t ) ∥ ≤ L ∥ x t -y t ∥ . Using Gronwall's inequality yields the result. Lemma A.2. Under Assumption 3.4, we have

PROOF 4.

for 0 ≤ t ≤ ϵ + t ≤ 1 .

PROOF 5. Direct application of Taylor approximation.

A.4 Visualization of Tasks

We provide a visualization of the 2D Maze Figure 7.

Demonstration (Single-task)

Start

Goal

Maze 1

Maze 2

Demonstration (Multi-task)

Maze 3

Maze 4

A.5 Planner for Maze2D task

We generate the demonstration data in Maze toy using planner similar to [54]. The planner devises a path in a maze environment by calculating waypoints between the start and target points. It begins by transforming the given continuous-state space into a discretized grid representation. Employing Q-learning, 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 an ordered manner to form a 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.

1

0

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 to 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 in our main paper. In Joint Training setting, we train both the variance prediction network and the policy function simultaneously in an 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 of our two-stage framework in balancing policy accuracy and uncertainty estimation. Separate training exhibited faster computational speed, making it the preferred choice once the policy function was robustly trained. Joint training required more computational resources and time.

A.7 Visualization of Exact Variance.

In the main paper, we 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, we first train a 1-Rectified Flow model for the task, then we can compute the exact variance by sampling:

For each states, we randomly sample 10 time steps ( N t = 10 ) and 10 noises ( N z = 10 ).

A.8 Visualization of Predicted Variance on Maze task.

We present the predicted variance by AdaFlow in Figure 8.

A.9 Additional Experimental Details.

Model Architectures. For the 1D toy example, we used a MLP constructed with 5 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 ∗ z T , where z T is sampled from a Gaussian distribution. This time feature is then concatenated

with the noise and condition ( x ) 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 as used in Push-T task. In the RoboMimic and LIBERO experiments, we used the Diffusion Policy-C architecture. To ensure a fair comparison across different methods, we maintained a consistent architecture for all methods in our experiments, except where specifically noted. Detailed parameters are available in Table 7.

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 a similar experimental pipeline. We just use a 4 layer MLP with SiLU activation for the variance prediction, with hidden dimension of 512, which is a very small network whose computational overhead can be neglected compared to the full model.

Implementation of Vairance Prediction Network. In the 1D 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, so 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 RoboMimic. 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 to three times longer 1 Given 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, we can consider standard Rectified Flow as a subset of our method, which can be recovered with specific choices of η 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, a 2-Rectified Flow is capable of producing a 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 Table 8, we compare the performance of 1-Rectified Flow, 2-Rectified Flow, and our method in the maze task. Furthermore, Figure 9 illustrates the trajectories produced by both standard Rectified Flow

1-Rectified Flow (NFE=1)

2-Rectified Flow (NFE=1)

1-Rectified Flow (NFE=5)

AdaFlow (NFE=1.12)

and our method. It's evident that the standard 1-Rectified Flow struggles to generate a diverse range of actions in a single step. In contrast, our method is able to produce diverse behaviors in nearly one step. This efficiency is attributed to our method's ability to estimate the variance across different states, identifying those that require multi-step generation.

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 or 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 to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are 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 a separate "Limitations" section in 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 on 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 a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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 a 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, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it 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 submissions 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 it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we 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 a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

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 a 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 to 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, hyperparameters, 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 in the core of the paper to a 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, or 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, confidence intervals, or 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, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 or 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 as 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 or failed experiments that didn't make it 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://neurips.cc/public/EthicsGuidelines ?

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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

generate deepfakes for disinformation. On the other hand, it is not needed to point out 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 is 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by 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 a 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 use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor 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 and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe paper lacks a discussion on the limitations of the method.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nNo, 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe \"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.\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\nMajor: \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\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}