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

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.

Strengths

- The idea of this paper is well-motivated and the investigated problem is practical. - The theme of this paper in speeding up flow-based generative models is timely and interesting. - Theoretical analysis is complete. - Empirical results show improvement over baseline and other compared methods. The visualized results are good to demonstrate the efficiency.

Weaknesses

- There are some minor typos in the paper. - The implementation codes are not provided. - The experimental data lack statistical significance. - The limitations of the method are not discussed.

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.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The paper lacks a discussion on the limitations of the method.

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

Strengths

- The paper proposes a simple yet efficient adaptive scheme to dramatically reduce the computational cost for inference in generative model policies. - The method is well-grounded in theoretical analysis. - The paper is generally well-written and clearly structured.

Weaknesses

- The environments tested in this work are standard but relatively low-dimensional. It's unclear how well the proposed method scales to high-dimensional tasks like humanoid control. - The experimental section could benefit from more detailed information about the setup and evaluation metrics.

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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

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.

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.

Strengths

The paper is concerned about the inference speed of learned imitator policies, which is an important problem that is sometimes ignored in academic work on IL, especially with the recent focus on diffusion models. The authors present a diverse set of examples that is somewhat suggestive of the efficacy of their techniques.

Weaknesses

The "theoretical analysis" (Line 59) presented is imprecise and holds only under extremely strong and unrealistic conditions. As an example of the imprecision, the conditioning in the definition of the conditional variance in line 147 is a bit confusing to me. Is the conditioning on $x_t$ intended to be conditional on the path up to time $t$? If not, then it is not obvious to me that we could expect that $\sigma^2$ can be zero in any setting where the gradient flow converges, due to the variance coming from $x_0$. If so, this should be formally defined. Furthermore, many of the claimed propositions, such as 3.1 and 3.2 are immediate. Proposition 3.3 has a non-rigorous proof that refers to a vague property of a rectified flow without stating precisely what this property is or citing the result. While in principle, any estimator of the conditional variance would be acceptable, the authors only present one that assumes the conditional distributions of the actions are Gaussian with no justification and no discussion of this assumption.

Questions

See weaknesses.

Rating

6

Confidence

3

Soundness

2

Presentation

2

Contribution

2

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.

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.

Strengths

* The paper shows an interesting idea - to use an estimation of the uncertainty in action prediction (the conditional variance) to dynamically adapt the step size of the flow model. If the model is certain in its action prediction, the step size will be large. Otherwise, the model performs many small prediction steps. * AdaFlow speeds-up the inference step of flow-based models, which is an important factor for many real-time applications. * The authors show the advantages of the proposed approach in various environments (4 tasks). * The paper is well-written and easy to follow and understand.

Weaknesses

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

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

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.

Thanks for your clarification

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.

Thank you for the answers

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 the additional experiments

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.

Thank you for the clarification

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.