AdaFlow: Imitation Learning with Variance-Adaptive Flow-Based Policies

Authors: xixi  hu
bo  liu
xingchao  liu
qiang  liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 00:40:22

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.

Main paper: 
 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] [4] [5] [6] .38th Conference on Neural Information Processing Systems (NeurIPS 2024). In the upper portion of the image, we use Diffusion Policy (DDIM) and AdaFlow to predict y given x, with deterministic y = 0 when x ≤ 0, and bimodal y = ±x when x > 0. Both DDIM and AdaFlow fit the demonstration data well. However, the simulated ODE trajectory learned by Diffusion-Policy with DDIM (red) is not straight no matter what x is. By contrast, the simulated ODE trajectory learned by AdaFlow with fixed step (blue) is a straight line when the prediction is deterministic (x ≤ 0), which means the generation can be exactly done by one-step Euler discretization. At the bottom, we show that AdaFlow can adaptively adjust the number of simulation steps based on the x value according to the estimated variance at x. The primary approach for IL is Behavioral Cloning (BC) [7] [8] [9] [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] [18] [19] [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] [22] [23] [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:• We proposed 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.• We offer a theoretical analysis of the overall error in action generation by AdaFlow, providing a bound that ensures precision and reliability.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 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. 
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] [31] [32] [33] , audio generation [34] , point cloud generation [35] [36] [37] [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] [40] [41] [42] [43] [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] [49] [50] [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.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. 
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.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,dz t = v θ (z t , t | s)dt, z 0 ∼ N (0, I).(1) 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,EQUATIONwhere x t is the linear interpolation between x 0 and x 1 = a:EQUATIONWe 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] .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,EQUATIONAfter 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 isEQUATIONwhose trajectories are straight lines pointing to z 1 and hence can be calculated exactly with one step of Euler step:z 1 = z 0 + v * (z 0 , 0 | s).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ϵ t = Clip η σ ϕ (z t , t | s), [ϵ min , 1 -t] .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 = ta + (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:EQUATION= E ∥a -x 0 -v * (x t , t | s)∥ 2 | x t = x, s .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:EQUATIONWe 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.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 :EQUATION)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 dz t = v * (z t , t | s)dt. Assume z t ∼ p t = p * t and p t+ϵt the distribution of z t+ϵt following (8). Then we haveW 2 (p * t+ϵt , p t+ϵt ) 2 ≤ ϵ 2 t E zt∼pt [σ 2 (z t , t | s)],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:EQUATIONwhere 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 writev * t (•) = v * (•, t | s). Assumption 3.4. Assume ∥v * t ∥ Lip ≤ L for t ∈ [0, 1], and the solutions ofdz t = v t (z t )dt has bounded second curvature ∥ zt ∥ ≤ 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 haveW 2 (p * 1 , p 1 ) ≤ C × N ada N max × ϵ min ,where C is a constant depending on M , M η and L.The idea is that the error is proportional to N ada Nmax , 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer 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). 
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. Behavior Diversity ✗ ✓ ✓ ✓ Fast Action Generation ✓ ✗ ✓ ✓ No Distillation / Reflow ✓ ✓ ✗ ✓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 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.EQUATIONWe 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 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 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.We valid how AdaFlow performs against baselines regarding the training and inference efficiency. In addition, we examine how critical the variance estimation network is.Predicted Variance Higher Training and Inference 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.  
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. 
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. 
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. • 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. 
noise: 
 Step 50Step 100 Noise Low-Variance State High-Variance StateStep 5Step 10State Variance1 Low HighFigure 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 a one-step action generator. For high-variance states, it employs more steps to ensure accurate action generation. This adaptive approach enables AdaFlow to achieve an average generation speed close to one step per task completion. 
 
 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 = ta + (1 -t)x 0 , note that a -x 0 = 1 1 -t (a -x t ).Therefore, ax 0 is deterministically decided by x t and s. This yieldsv * (x, t | s) = E[a -x 0 | x t = x] = 1 1 -t (a -x).Therefore, we havedz t = v * (z t , t | s) = 1 1 -t (a -z t )dt.Solving ODE this yieldsz t = ta + (1 -t)z 0 = (1 -t)v * (z 0 , 0 | s).Differentiating it also yields z t = (a -z 0 )dt. We also have σ 2 (x, t | s) again because a -x 0 is deterministic given x t and s:σ 2 (x, t | s) = var(a -x 0 | x t = x, s) = 0.A.2 Proof of Proposition 3.3 PROOF 2. Following the property of rectified flow, the distribution of x 1 = ta + (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 havez t+ϵt = x t + ϵ t v * (z t , t | s) and x t+ϵt = x t + ϵ t (a -x 0 ). We have W 2 (p * t+ϵt , p t+ϵt ) 2 ≤ E ∥z t+ϵt -x t+ϵt ∥ 2 2 = E E ∥z t+ϵt -x t+ϵt ∥ 2 2 | x t = E E ∥ϵ t v * (z t , t | s) -ϵ t (a -x 0 )∥ 2 2 | x t = ϵ 2 t E zt∼pt [σ 2 (z t , t | s)].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 ti 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 dz ti t = v t (z ti t )dt 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 ODEdz * t = v * t (z * t )dt. From Lemma A.1, we have z ti-1 1 -z ti 1 ≤ exp(L(1 -t i )) z ti ti -z tt-1 ti.Let p ti t be the distribution of z ti t . Then we have p 1 t = p t and p 0 t = p * t . ThenW 2 (p ti-1 1 , p ti 1 ) ≤ E z ti-1 1 -z ti 1 2 1/2 = exp(L(1 -t i ))E z ti-1 ti -z ti ti 2 1/2 = exp(L(1 -t i )) max(η, ϵ 2 min M/2) = Cϵ 2 min exp(-Lt i ),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,W 2 (p * 1 , p 1 ) = N ada i=1 W 2 (p ti-1 1 , p ti 1 ) ≤ N ada i=1 Cϵ 2 min exp(-Lt i ) ≤ C × N ada N max × ϵ min , 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 dx t = v t (x t )dt and dy t = v t (y t )dt starting from x 0 , y 0 , respectively. We haveEQUATION) PROOF 4. d dt ∥x t -y t ∥ 2 = 2(x t -y t ) ⊤ (v t (x t ) -v t (y t )) ≤ 2L ∥x t -y t ∥ 2 ,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∥x t+ϵ -(x t + ϵv t (x t ))∥ ≤ ϵ 2 M 2 , for 0 ≤ t ≤ ϵ + t ≤ 1.PROOF 5. Direct application of Taylor approximation.We provide a visualization of the 2D Maze Figure 7 . Start GoalWe 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. 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:EQUATION)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 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 Flow1-Rectified Flow (NFE=1) 2-Rectified Flow (NFE=1) 1-Rectified Flow (NFE=5)AdaFlow (NFE=1.12)Figure 9 : Generated trajectories. We visualize the 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 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 ChecklistSee this link 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: Major: 
* 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.\
\
Overall, 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.\
\
In 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).



* 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).

Minor mistakes: 
* Figure 1 is not referenced in the paper and also the caption does not include an explanation of its subfigures. 
* Line 175 - larger instead of large
* Table 5 - Maze 2 instead of Maze 1
questions: questions: * In line 174 - should the inequality for epsilon be with $\sqrt{\eta}$ instead of $\eta$?
* 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?
* 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.   
* Can the adaptive step size be implemented to generate images? If yes, what are the complications that may arise?

I am willing to raise my score if the authors address all of my concerns.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: questions: See weaknesses.
limitations: 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.
paper_summary: paper_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
strength_weakness: strength_weakness: - 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.
- The experimental section could benefit from more detailed information about the setup and evaluation metrics.
questions: 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?
- 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.
- How does AdaFlow compare with classical adaptive ODE solvers?
- 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?
- What is the robustness of AdaFlow to potential errors in the variance estimation network?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - There are some minor typos in the paper.
- The implementation codes are not provided.
- The experimental data lack statistical significance.
- The limitations of the method are not discussed.
questions: 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?
- 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]?
- 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."
- 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.
- What does "Hz" mean in Table 1?
- In formula (7), the log likelihood should be $\log(\sigma^2)$ instead of $\log(\sigma)$.
- Please carefully check and correct the formatting of references, e.g., [36] lacks author information, and references like [15], [32], [34] lack dates.

[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: limitations: The paper lacks a discussion on the limitations of the method.

Comments: title: Reply to Reviewer's Comment
Thank you for the clarification
Thanks for the additional experiments
Thank you for the answers
Thanks for your clarification
comment: We thank the reviewer for clarifying the question and raising the rating. Regarding the assumption that $a - x_0 - v$ is Gaussian, here is the explanation.

For simplicity, we temporarily ignore the dependency on state $s_t$, then:

$$
L(\sigma) = \mathbb{E}\bigg[\int_{0}^1 \bigg(\frac{||(A - X_0) - v(X_t, t )||^2}{2\sigma(X_t, t)^2} + \log \sigma(X_t, t) \bigg) dt \bigg]
$$

First, assume $v(x_t, t)$ is learned well such that $v(x_t, t) = \mathbb{E}\big[ A - X_0 \mid X_t = x_t \big].$ This is guaranteed by the property of Rectified Flow.

Next, take $\frac{dL}{d\sigma(x_t, t)}$ and set it to zero, we have:

$$
\frac{dL}{d\sigma(X_t, t)}|_{X_t = x_t} = \mathbb{E}\bigg[||(A - X_0)- v(X_t, t)||^2 \frac{-1}{\sigma(X_t, t)^3} + \frac{1}{\sigma(X_t, t)} \mid X_t = x_t\bigg] = 0
$$

Hence:

$$
\sigma(x_t, t)^2 = || (A - X_0) - v(X_t, t) \mid X_t = x_t ||^2 = \text{Var}(A - X_0 \mid X_t = x_t)
$$

As a result, note that though the particular objective of $L(\sigma)$ looks like the negative log likelihood of a Gaussian, **the closed-form of it indicates that the minimizer $\sigma^2$ corresponds exactly to the variance, no matter what true distribution it should be.** In other words, even if the $A - X_0 - v$ does not follow a Gaussian, minimizing Equation 7 (the objective of the $\sigma$ network) will still give the right variance estimation.
Thank you for the clarification regarding line 147.  On the subject of Proposition 3.3, I agree that the conclusion is correct, but I maintain that the proof as written is not rigorous and should be cleaned up.  I am happy to go along with the other reviewers on the subject of the experiments and raise my score accordingly due to the experimental results, but I still think that the theory is somewhat lacking. 

To clarify my point regarding the assumption of Gaussianity, note that equation (7), which you use to estimate the conditional variance is only consistent under the assumption that $a - x_0 - v$ is Gaussian.  While I believe that this may hold in practice due to universality properties of continuous stochastic processes, I think that this assumption should be more clearly stated and discussed.
Thanks for the rebuttal. Some of my concerns have been solved. 

I appreciate the additional comparison with DPM solvers.  Regarding the higher-dimensional tasks, I don't think the LIBERO dataset is a convincing example. Correct me if I am wrong, in LIBERO dataset, although images are part of the inputs, they are only served as a condition. The output is still the low-dimensional action space. Also, given your reply to Reviewer df3a about AdaFlow on image generation, it seems like AdaFlow will offer a marginal improvement on high-dimensional tasks. 

Regarding the clip of $\epsilon_t$, it cannot prevent the case that the true variance is high, but the estimation is low. 

In the light of this, I would keep my score.
I thank the authors for their response. The authors address most of my concerns - The additional explanation regarding the training time of the variance network is important and should be added to the revised version. Also, the comparisons to DPM and DPM++  solvers are important and strengthen the paper's contribution.


Overall, considering the additional explanations and experiments provided in the rebuttal period and all reviewers' answers, I decided to raise my score to "accept".
Thanks for your rebuttal. My concerns have been resolved, except for the third point in Weaknesses. The standard deviation or error bars of the experimental results need to be clearly provided to avoid the randomness. Anyway, I think this is an interesting job for the field of imitation learning. I decided to improve my score.