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

Thank you for your valuable feedback, which has helped us improve the clarity and comprehensiveness of our work.

Complexity of expert behavior and variance

The complexity of expert behavior and the variance are intertwined. Since we perform offline learning with demonstration data, the model learns a high variance when the expert’s behavior becomes diverse at a given state, and vice versa. So even if the expert behavior is very complicated (like peg in the hole), but the action itself is relatively deterministic (there is a consistent way of solving the underlying task), then AdaFlow will still choose a large step size, because the model is more certain about which action to take, so it performs more like in the BC case. In contrast, even if it is a simple task like lifting an object, if there exists many different ways of lifting, then AdaFlow will still choose a small step size, until the lifting behavior is determined (the conditional variance drops accordingly).

Comparison with Reflow

Yes, Reflow can achieve comparable success rates to AdaFlow. However, Reflow requires an additional complex and time-consuming stage after training the original flow. AdaFlow is much more efficient than Reflow. This is because:

• AdaFlow doesn’t have to be trained in two separate stages, it can also learn action generation and variance estimation together. As shown in the following table (originally in Appendix A.6), training the flow and the variance networks separately or jointly results in similar performance. In practice, we chose to train them separately because separate training in fact results in even better training efficiency (faster convergence), this is because the objective of the variance network is stationary. In addition, training the variance network is way faster than training the original flow, hence Adaflow introduces almost no overhead to the 1-Rectified Flow (without reflow). Compared to AdaFlow, reflow and consistency distillation always require a stage to distill the ODE，which itself alone is more expensive than the original training phase.

  	Maze 1	Maze 2	Maze 3	Maze 4
  AdaFlow (Separate)	0.98	1.00	0.96	0.86
  AdaFlow (Joint)	1.00	1.00	0.96	0.88

• Reflow is more complicated and computationally expensive than AdaFlow. Reflow operates in two steps: first, it generates pseudo data, requiring simulating the full trajectory, i.e., 100 model inferences per pseudo sample, and the reflow dataset usually needs to be as large as the original dataset. Second, it trains the new flow model with a reflow target, which might introduce errors. AdaFlow doesn’t need this expensive data generation. If we train the variance estimation part separately (it can also be trained jointly), it only involves training a few linear layers on top of the original rectified flow model. In practice, we compared the total training time for the reflow stage and AdaFlow’s variance estimation training stage. On the maze task, we observed that reflow takes 7.5x longer than AdaFlow in terms of wall clock time.

• It is possible to apply a distillation stage after reflow. In image and video generation tasks, distillation is often necessary because the ODE might not be straight even after reflow. However, in our robotics tasks, reflow has already learned a very straight trajectory. We observed no further improvements when using more than one inference step for reflow. Therefore, the potential improvement from distillation is very minimal. Additionally, distillation requires another training stage, which adds complexity and inconvenience.

Measure behavior richness

We appreciate the suggestion to use NMI or ENT as a measure of behavioral richness. However, it's not feasible to apply them directly to the generated trajectories because each trajectory has a different length. Clustering methods like NMI typically require fixed-length data, which isn't the case with our generated trajectories. We are open to exploring alternative metrics that could better capture the diversity and richness of behaviors in future work. If you have any suggestions or additional metrics that you believe could be useful, we would be happy to consider them.

Typos

There’s a typo in Table 1, which we will correct in the revised version. We will correct other typos mentioned by the reviewer in the paper and will reflect these changes in the revised version.

Please let us know if you have any further questions or concerns. We hope that our rebuttal addresses all your points thoroughly, and we would greatly appreciate your reconsideration of the score. Thank you!

We believe the reviewer might have some misunderstandings, and we will ensure to further clarify our statements in the updated version of the paper.

• In Line 147, $x_t$ is a random variable that follows the marginalized probability at time $t$. Thus, $x_t$ is not conditional on the path up to time $t$ (similar to the original Rectified Flow paper). The conditional variance $\text{Var}(a - x_0 | x_t = x, s)$ can indeed be zero, even though $x_0$ is random. This is because the variance is conditional on $x_t = x$. For instance, in a 1-dimensional example where $a = 1$, $x_0$ follows a Gaussian distribution $N(0, 1)$, and we define $x_t = (1 - t) x_0 + t$, as in Rectified Flow. In this case, given $x_t = x$, we can deterministically solve $(a - x_0)$ because $x_0 = \frac{x_t - t}{1 - t}$ and $a$ is a constant 1. In fact, in this simple 1D example, the flow is exactly straight because $a$ is deterministic, which is what Proposition 3.2 is stating.

• Regarding Proposition 3.3, the statement is accurate and correct. However, we will ensure that the statement is made more rigorously with a clear presentation of the assumptions and properties from earlier works.

• The property mentioned is that $z_t$ and $x_t$ share the same law (i.e., for any $t$, their distributions are the same), which is formally stated as Theorem 3.3 in the Rectified Flow paper.

• We respectfully note that we do not assume the conditional distributions of the actions are Gaussian. The action can follow any distribution, similar to the image distribution in image diffusion applications. Only the initial noise $x_0$ follows a Gaussian distribution. We would appreciate it if the reviewer could clarify what specific assumption or detail they are referring to regarding the Gaussian distribution.

Thank you for your valuable feedback, which has helped us improve the clarity and comprehensiveness of our work.

Possibility of extending to higher-dimensional robotics tasks

We agree that applying AdaFlow to more complex and high-dimensional tasks like humanoid control would be interesting. Due to the time constraints of the rebuttal period, we couldn't include such experiments. However, on the LIBERO dataset, our input is video data plus the gripper configuration of the robot hand, which is one of the most high-dimensional tasks in robot manipulation. In future work, we plan to explore using AdaFlow on humanoid tasks. We are confident that training an ODE on higher-dimensional spaces is feasible, as ODE training has been successful in image generation tasks (as in Stable Diffusion 3)

Detailed experiment setup

We will update the paper to include more detailed information about the experimental setup and evaluation metrics in Appendix A.9 Additional Experimental Details. Additionally, we will provide the codebase to assist with reproducing the results.

Correction of typos

We have corrected the typos in the paper to ensure consistency and will reflect these changes in the revised version.

Seeds and runs

Due to the limited time during the rebuttal period, we couldn't complete the experiments with multiple seeds and runs. However, the partial results we obtain indicate that the performance of AdaFlow, diffusion policy, and other baselines are consistent with what we presented in the paper. We will include the results with multiple seeds in the revision to provide a comprehensive evaluation.

Comparison with classical adaptive ODE solvers

We conducted experiments to compare AdaFlow with recent advanced and faster solvers like DPM-solver and DPM-solver++. As shown in the table below, we compared AdaFlow with DPM and DPM++ across multiple tasks on LIBERO. The results demonstrate that AdaFlow achieves higher success rates even with fewer NFE:

Robustness to potential errors in the variance estimation network

We acknowledge that the variance estimation network may have errors. As shown in this figure, the predicted variance does not perfectly match the ground truth state variance. This estimation error can potentially influence the generation results.

However, note that in Algorithm 1, $\epsilon_t$ is clipped between $\epsilon_\text{min}$ and $(1 - t)$, so in theory, AdaFlow’s worst case performance (with arbitrarily bad variance estimation) will be between 1-step RF and $1 / \epsilon_\text{min}$-step RF.

Empirically, we compare AdaFlow against RF with 20 inference steps on the LIBERO task suites, the results are shown in the following:

As shown, AdaFlow's performance is comparable to Rectified Flow with NFE=20. This demonstrates that the introduction of the variance estimation process does not significantly impact performance. AdaFlow remains effective and efficient.

We thank Reviewer df3a for the thoughtful questions and suggestions.

1. What is the advantage of AdaFlow over Rectified Flow with reflow?

• AdaFlow can learn action generation and variance estimation together, as shown in Appendix A.6, where joint and separate training yield similar performance. We opted for separate training due to better efficiency and faster convergence, given the stationary objective of the variance network. Training the variance network is much quicker than training the original flow, adding minimal overhead to the 1-Rectified Flow. In contrast, reflow and consistency distillation always require an expensive ODE distillation stage.

• Reflow is more complex and computationally expensive than AdaFlow. Reflow requires generating pseudo data by simulating the full trajectory (100 model inferences per pseudo sample) and needs a reflow dataset as large as the original. Then, it trains a new flow model with a reflow target, which can introduce errors. AdaFlow avoids this costly data generation. Training the variance estimation part separately (or jointly) involves only a few linear layers on the original rectified flow model. In practice, on the maze task, reflow takes 7.5 times longer than AdaFlow in terms of wall clock time.

2. Why does reflow lead to inferior generation quality with straightened ODE?

• Reflow estimates a new mapping ODE from noise to data, introducing errors. The first error comes from time-discretization during numerical simulation, and the second from training on pseudo data pairs. AdaFlow avoids these issues by using the original ODE.

• In Tables 2 and 4, Reflow's rare outperformance of AdaFlow isn't due to variance prediction errors but a balance of efficiency and success rate. With more inference steps, AdaFlow outperforms Reflow, demonstrating a better ODE. The following tables compare Reflow and AdaFlow with different inference steps (NFE) on LIBERO tasks.

  	Task 1	Task 2	Task 3	Task 4	Task 5	Task 6	Average
  Reflow	0.90	0.82	0.98	0.82	0.82	0.96	0.88
  AdaFlow (NFE=1.27)	0.98	0.80	0.98	0.82	0.90	0.96	0.91
  AdaFlow (NFE=2.91)	0.96	0.86	1.00	0.82	0.90	0.98	0.92

3. Comparison to DPM solver and DPM ++ solver

Thank you for suggesting a comparison with advanced solvers like DPM-solver and DPM-solver++. We conducted additional experiments, and as shown in the table below, AdaFlow achieves higher success rates even with fewer NFE across multiple LIBERO tasks.

4. Difference to RK45 solver

The RK45 solver adjusts step size for accuracy by using the difference between its 4th and 5th order estimates to calculate an error estimate, which then adjusts the next step size.

AdaFlow and RK45 are orthogonal methods for step size estimation:

• RK45 increases numerical simulation accuracy.
• AdaFlow leverages the straightness of the learned flow.

Thus, they can be used together. The theoretical lowest NFE for RK45 is 6, as it calculates 6 intermediate values per step, while AdaFlow needs close to one step.

5. Use other non-post-training methods to estimate the state variance

We explored estimating state variance without training a variance estimation network. We trained only the rectified flow model and estimated state variance by sampling. For a given state, we sampled $N_z$ noises and $N_t$ timesteps, and then estimate the variance of each state following Equation 6.

$\frac{1}{N_t} \frac{1}{N_z} \sum_t\sum_{z_0} \mathbb{E}[||y - z_0 - v(z_t, t; x)||^2]$

Here $y$ is derived using the Euler sampler, and $z_0$ is random Gaussian noise. This method estimates state variance accurately if both $N_z$ and $N_t$ is large enough. However, it involves sampling the full flow trajectory for many different noises. The figure below shows how many samples are needed to achieve a precise estimation of state variance: Figure Link. In general, an NFE > 10 is required to estimate the variance well. Compared to this method, AdaFlow is much faster for inference, only requiring one model inference.

6. Usage in generating images

Great point! We've also tried using AdaFlow for image generation. Applying the learned variance network at inference does speed up the process. However, the speedup is limited because the learned flow is not straight due to the high dimensionality and variability in image generation.

In robotics tasks, many actions are deterministic given the states, unlike in image generation. For text-conditioned images, even detailed prompts can result in many possible images, making the generation trajectory non-deterministic. Thus, AdaFlow is better suited for robotics, where deterministic generations significantly reduce inference costs.

7. Typos We will correct the typos mentioned by the reviewer and will reflect these changes in the revised version.

Please let us know if you have any further questions or concerns. We hope that our rebuttal addresses all your points thoroughly, and we would greatly appreciate your reconsideration of the score. Thank you!