Paper: RF-POLICY: Rectified Flows are Adaptive Decision Makers (ICLR-2024)

Reviewer comment (Experiments):
Some claims of the paper are not adequately supported with evidence. For example, the experiments evaluate the method on three dimensions only on two benchmark datasets (the third is in the appendix and performance is comparable), so the last claim in the abstract may be doubted. Another example is that Assumption 1 is too strong and there is no evidence nor data distribution visualization to support it, thus the use of the proposed method with real data and non-expert data (e.g., low-return trajectories) is not convincing based on the limited results shown in this paper. However, the focus of this work is to improve the computational efficiency of diffusion-based policy generators, thus more experiments on other datasets as used in Chi (2023) should be conducted, especially the experiments on real-world robot benchmark if possible.

Author response:
**Q1. Assumption 1 and Real Data Application:**

Regarding Assumption 1, we acknowledge the need for empirical support. To this end, we have conducted further analysis to provide data distribution visualizations that substantiate our assumption. These additional results will bolster the credibility of applying RF-POLICY to real and non-expert data.

To substantiate this assumption, we have conducted additional analyses. 
We take advantage of the variance predictions made by RF-POLICY across different states within a robot's state space. Figure 8 in our paper ([link](https://anonymous.4open.science/r/RF-POLICY-184B/libero_variance.png)) presents these predictions, offering a visualization of the diversity in potential actions given a robot's specific state. As illustrated in the figure, for the pick and place task, there is an observable increase in action diversity as the robot's end effector gets closer to the target object. This pattern suggests that while most states exhibit low action variance, implying a uni-modal distribution of potential actions, a subset of states, particularly those in proximity to the object interaction, display a significant increase in variance, reflecting a multi-modal distribution of actions.

**Q2. Comparisons with Chi et. al.:**

We acknowledge the importance of real-world robot experimentation to validate our findings, however, at the current moment the authors do not have real robots to conduct such experiments. As a compensation, we conduct more experiments in simulation (Please see the general response).

Regarding benchmarks from Chi et al.'s work, we have indeed included results from the robomimic benchmark in the appendix of our initial submission (Section 1 in Appendix). Notably, we observed that the success rates for BC, Diffusion Policy, and RF-POLICY were all exceedingly high, consistently above 0.95, indicating a ceiling effect that limits the differentiation between methods. To address this, we shifted our focus to the LIBERO benchmark, which presents more challenging and discriminative tasks.

We plan to undertake a comprehensive set of experiments, including those on RoboMimic, and will incorporate these results in an updated version of our work.

Reviewer comment (Experiments):
The proposed method is compared with only BC and Diffusion Policy, though the baselines are representative in either field. Performance improvements over other diffusion-based polices are missing, such as Diffusion-QL (Wang et al., 2023), Decision Diffuser (Ajay et al., 2023), Difusion BC (Pearce et al., 2023), etc, so the significance of this paper is doubted.

Author response:
**Q3: Comparison with other offline RL learning methods:**

Our work focuses on offline imitation learning, not offline reinforcement learning, which means the rewards are not present. This is a more practical assumption at least in robotics, since learning complex manipulation/navigation behaviors often requires complex reward design, while providing demonstrations is much simpler. As a result, our method is not directly comparable against methods like Diffusion-QL and Decision Diffuser. Regarding Diffusion BC, it uses DDPM to train a diffusion model for action prediction, this objective is the same as Diffusion Policy we compared with. 

In response to the review, we have expanded our comparative analysis to include additional relevant baselines within the scope of imitation learning (See general response). These additional comparisons are better aligned with our approach and further validate the efficacy of RF-POLICY in the intended setting.

Reviewer comment (Presentation):
Description of some figures and experimental results is confusing and need further clarity.

Author response:
We thank the reviewer for pointing out the typos in Figures and tables. We have fixed the typos and revise the descriptions of figures and experimental results for clarity. Please see the updated paper for these change. Considering the time constraints during rebuttal, we will include the average scores across multiple seeds to ensure that the results reflect a robust evaluation of our method.

Reviewer comment (Theory):
The paper lacks empirical or mathematical validation for Assumption 1. The authors posit that most demonstration datasets for robotic tasks exhibit deterministic behavior (or uni-modal states), yet they fail to support this claim with experimental evidence.

Author response:
**Q1. Validation of Assumption 1:**

To substantiate this assumption, we have conducted additional analyses. 

We take advantage of the variance predictions made by RF-POLICY across different states within a robot's state space. Figure 8 ([link](https://anonymous.4open.science/r/RF-POLICY-184B/libero_variance.png)) presents these predictions, offering a visualization of the diversity in potential actions given a robot's specific state. As illustrated in the figure, for the pick and place task, there is an observable increase in action diversity as the robot's end effector gets closer to the target object. This pattern suggests that while most states exhibit low action variance, implying a uni-modal distribution of potential actions, a subset of states, particularly those in proximity to the object interaction, display a significant increase in variance, reflecting a multi-modal distribution of actions.

Reviewer comment (Experiments):
The study omits a comparative analysis with established offline RL baselines such as CQL and BCQ, as well as other diffusion-based methodologies like Diffuser (Janner et al., 2022) and Diffusion-QL (Wang et al., 2022).

Author response:
**Q3: Comparison with other offline imitation learning methods:**

Our work focuses on offline imitation learning, not offline reinforcement learning, which means the rewards are not present. This is a more practical assumption at least in robotics, since learning complex manipulation/navigation behaviors often requires complex reward design, while providing demonstrations is much simpler. As a result, our method is not directly comparable against methods like CQL and BCQ.

In response to the review, we have expanded our comparative analysis to include additional relevant baselines within the scope of imitation learning (See general response). These additional comparisons are better aligned with our approach and further validate the efficacy of RF-POLICY in the intended setting.

Reviewer comment (Theory):
How do linear flow models (linear ODE models), like RF-Policy, accurately encapsulate complex behaviors? Most existing methods have relied on non-linear SDEs, specifically DDPM, for policy estimation, yet this study utilizes a linear model. What rationale is provided for the superiority of this linear approach over its predecessors? (Related to equation 3)

Author response:
**Q4: How does flow model capture complex behaviors?**

RF wants to construct ODEs that follow straight trajectories, which is very different from assuming a linear model. Constructing straight ODE is a non-trivial task, requires quite a lot of non-linearity to achieve: the ODE has to control the move of the particle carefully to ensure the straight trajectory. Mathematically, the velocity field of straight ODEs is characterized by a nonlinear PDE called burger’s equation[1]:

$\partial v/ \partial t + (\partial v/ \partial z) v = 0$, where $v = v(Z_t, t)$ is the velocity field.

This is a challenging and highly non-linear problem to solve. RF seeks for straightness because it is easier for simulate at inference time (more straight => fewer steps).

[1] See Liu, et. al., Flow Straight and Fast. 2023

Reviewer comment (Presentation):
Figure 1 is intended to demonstrate that, unlike DDIM (an extension of DDPM), RF-Policy generates straight lines in deterministic areas (x < 0). However, I cannot see distinguishable difference between the red (DDIM) and blue (RF-Policy) lines in the figure. It could be considered a potential weakness of the paper.

Author response:
**Q5. Clarifying Trajectories in Figure 1:**

Figure 1 presents a 1D example with mappings that highlight the behavior of both DDIM and RF-Policy. Although both methods fit the demonstration data well, the red line (DDIM) demonstrates a non-linear trajectory, indicating that multiple inference steps are necessary due to the nature of SDEs. In contrast, the blue line (RF-Policy) presents straight trajectories for deterministic cases (x ≤ 0), validating RF-Policy's ability to efficiently generate actions in a single step through one-step Euler approximation ​

Reviewer comment (Theory):
How does the variance prediction network determine whether a state is uni-modal or multi-modal? It is trained to estimate state variance using an offline dataset, encompassing both epistemic and aleatoric uncertainties. Given that the distinction between uni-modal and multi-modal states pertains to aleatoric uncertainty, how does the model address epistemic uncertainty?

Author response:
**Q6:How does the variance prediction network determine whether a state is uni-modal or multi-modal?**

We thank the reviewer for bringing up this question. We want to first emphasize that the multi-modality here is a bit different from the term uncertainty. It is possible that the data has no uncertainty at all, but we want to learn a multi-modal behavior (at certain state, the Q(s, a1) = Q(s, a2) >= Q(s, other_a), so both a1 and a2 are optimal actions). Therefore, at this moment, RF-policy, like other offline imitation learning methods (e.g., BC and Diffusion-policy) are not designed to address the epistemic uncertainty. The epistemic uncertainty could be addressed by either adding more data or extending RF-policy to estimate the (s, a) distribution. Both solutions are also applicable to diffusion policy, but we argue that either case only worsen the computation cost. On the contrary, RF-policy is more training/inference efficient, making it a better candidate to scale up to more data or more complex modeling (like world-model learning).

Reviewer comment (Experiments):
The argument would be materially stronger if the method was demonstrated on a variety of more realistic domains than the toy domains (e.g., maze) used in the paper. For example, Atari-100k seeded with expert trajectories could be used, which should not require industrial levels of compute.

Author response:
**Q1: Addressing More Realistic Domains:**

We appreciate the reviewer's suggestion to test RF-POLICY on more complex domains, such as the Atari-100k benchmark. However, Atari games, including benchmarks like Atari-100k, use a discrete action space. This poses a challenge for diffusion models, and it may require additional methods to approximate the continuous processes that diffusion models use.

In light of this, and to demonstrate the versatility of RF-POLICY in more sophisticated and realistic scenarios, we have extended our experiments to include a 'pick and place' task. This task incorporates a set of demonstrations showcasing different methods of picking up a bowl. We believe this task mirrors real-world applications more closely, as it involves learning multiple viable solutions for a single task – a common requirement in practical settings. This extension of our experiments serves to better validate the model's effectiveness in complex, real-world applications where diverse solution strategies are essential. The following figure shows two ways of picking up the bowl (See this [link](https://anonymous.4open.science/r/RF-POLICY-184B/diverse_libero.png) or Figure 7 in our paper). The table shows the comparison between bc, standard RF, Rectified Flow and RF-POLICY (See this [link](https://anonymous.4open.science/r/RF-POLICY-184B/diverse_libero.png) or Table 8 in our paper). 

|            | SR     | DS   | NFE  |
|------------|--------|------|------|
| BC         | 0.92   | 0.90 | 1.00 |
| DDIM       | 0.96   | 0.97 | 10.00|
| RF-POLICY  | 0.96   | 0.95 | 1.37 |

Table 8: Comparison of Behavioral Cloning (BC), Diffusion Policy (DDIM), and RF-POLICY in Pick and Place tasks. The table includes performance metrics such as Success Rate (SR), Diversity Score (DS), and Number of Function Evaluations (NFE) for each model.

Reviewer comment (Experiments):
The work lacks an ablation study and exhaustive comparison with diffusion samplers specifically optimized for the low-NFE regime. Relevant baselines to compare against would include UniPC, Heun and others for pure samplers, as well as broader approaches like Consistency Models or TRACT.

Author response:
**Q2: Conducting Ablation Studies and Exhaustive Comparisons:**

We pick step=20 in the original paper because that is the number of steps used by Diffusion Policy. We appreciate the reviewer's emphasis on the necessity of a thorough comparison with diffusion samplers optimized for the low-NFE regime. To address this, we have expanded our comparative analysis to include DDIM and standard 1-RF at varying levels of NFEs. The following table presents this expanded evaluation, offering a clearer perspective on the relative performance trade-offs between these methods when operating under constrained NFEs (See this [link](https://anonymous.4open.science/r/RF-POLICY-184B/table7_low_nfes.png) or Table 7 in our paper).

Reviewer comment (Experiments):
The evaluation in Figure 5 appears to only use DDIM with 20 steps as the baseline. This choice seems arbitrary, and it is unclear what relative performance loss would occur with DDIM10, DDIM5, or other samplers. Without such comparisons, claims about NFE speedup are minimized.

Author response:
|               | Maze 1      |       |       |       | Maze 2      |       |       |       | Maze 3      |       |       |       | Maze 4      |       |       |       | Average     |       |       |       |
|---------------|-------------|-------|-------|-------|-------------|-------|-------|-------|-------------|-------|-------|-------|-------------|-------|-------|-------|-------------|-------|-------|-------|
|               | SR          | DS    | NFE   | LEI   | SR          | DS    | NFE   | LEI   | SR          | DS    | NFE   | LEI   | SR          | DS    | NFE   | LEI   | SR          | DS    | NFE   | LEI   |
| DDIM (step=1) | 0.12        | 1.00  | 1.00  | 0.07  | 0.04        | 1.00  | 1.00  | 0.03  | 0.04        | 1.00  | 1.00  | 0.01  | 0.02        | 1.00  | 1.00  | 0.01  | 0.06        | 1.00  | 1.00  | 0.03  |
| DDIM (step=3) | 0.88        | 1.00  | 3.00  | 0.64  | 0.78        | 1.00  | 3.00  | 0.54  | 0.90        | 1.00  | 3.00  | 0.55  | 0.84        | 1.00  | 3.00  | 0.52  | 0.85        | 1.00  | 3.00  | 0.56  |
| DDIM (step=5) | 0.88        | 1.00  | 5.00  | 0.64  | 0.76        | 1.00  | 5.00  | 0.53  | 0.98        | 1.00  | 5.00  | 0.65  | 0.96        | 1.00  | 5.00  | 0.69  | 0.90        | 1.00  | 5.00  | 0.63  |
| DDIM (step=10)| 0.92        | 1.00  | 10.00 | 0.66  | 0.84        | 1.00  | 10.00 | 0.55  | 1.00        | 1.00  | 10.00 | 0.64  | 0.94        | 1.00  | 10.00 | 0.68  | 0.92        | 1.00  | 10.00 | 0.63  |
| DDIM (step=20)| 0.92        | 1.00  | 20.00 | 0.65  | 0.78        | 1.00  | 20.00 | 0.56  | 0.98        | 1.00  | 20.00 | 0.60  | 0.92        | 1.00  | 20.00 | 0.69  | 0.90        | 1.00  | 20.00 | 0.63  |
| 1-RF (step=1) | 1.00        | 0.95  | 1.00  | 0.94  | 1.00        | 0.93  | 1.00  | 0.92  | 1.00        | 0.26  | 1.00  | 0.84  | 0.96        | 0.59  | 1.00  | 0.75  | 0.99        | 0.68  | 1.00  | 0.86  |
| 1-RF (step=2) | 1.00        | 0.99  | 2.00  | 0.85  | 0.98        | 0.96  | 2.00  | 0.85  | 1.00        | 0.78  | 2.00  | 0.82  | 0.92        | 0.87  | 2.00  | 0.76  | 0.97        | 0.90  | 2.00  | 0.82  |
| 1-RF (step=3) | 1.00        | 0.99  | 3.00  | 0.89  | 0.96        | 0.99  | 3.00  | 0.75  | 1.00        | 0.98  | 3.00  | 0.83  | 0.96        | 0.99  | 3.00  | 0.79  | 0.98        | 0.99  | 3.00  | 0.82  |
| 1-RF (step=5) | 1.00        | 1.00  | 5.00  | 0.93  | 0.92        | 1.00  | 5.00  | 0.75  | 1.00        | 0.99  | 5.00  | 0.83  | 0.92        | 1.00  | 5.00  | 0.84  | 0.96        | 1.00  | 5.00  | 0.84  |
| 1-RF (step=10)| 1.00        | 1.00  | 10.00 | 0.94  | 0.90        | 1.00  | 10.00 | 0.69  | 1.00        | 1.00  | 10.00 | 0.88  | 1.00        | 1.00  | 10.00 | 0.88  | 0.97        | 1.00  | 10.00 | 0.85  |
| 1-RF (step=20)| 1.00        | 1.00  | 20.00 | 0.92  | 0.82        | 1.00  | 20.00 | 0.68  | 1.00        | 1.00  | 20.00 | 0.84  | 0.94        | 1.00  | 20.00 | 0.86  | 0.94        | 1.00  | 20.00 | 0.83  |
| RF-POLICY     | 1.00        | 0.98  | 1.11  | 0.91  | 1.00        | 0.79  | 1.03  | 0.83  | 1.00        | 0.96  | 1.99  | 0.78  | 0.98        | 0.93  | 1.85  | 0.59  | 0.99        | 0.91  | 1.50  | 0.78  |
 
Caption: Comparative Performance of RF-POLICY, DDIM, and 1-RF Across Different Low-NFE Settings. The table showcases the Success Rate (SR), Diversity Score (DS), Number of Function Evaluations (NFE), and Learning Efficiency Index (LEI) achieved by each method under a range of NFE constraints.

In addition, We have expanded the application of RF-POLICY to serve as an adaptive step size ODE solver for unconditional image generation, integrating our variance estimation technique with a U-Net-based DDPM++ architecture for image synthesis on the CIFAR-10 dataset. Our performance evaluation used Frechet Inception Distance (FID) as the metric.

Our analysis reveals that RF-POLICY allows for larger step sizes and increased generation speeds in image synthesis when the ODE trajectory demonstrates linearity. In addition, it performs better than RK-45 sampler. The result can be seen here ([link](https://anonymous.4open.science/r/RF-POLICY-184B/cifar10.png) or Figure 6 in our paper).

Reviewer comment (Novelty):
It is not evident that Rectified Flow is the unique or best solution for low-NFE acceleration in offline IL. Other approaches, such as Consistency Models, could also potentially claim similar advantages in this setting.

Author response:
**Q3. Exploring Alternatives to DDIM:**

It is important to note that RF-POLICY is not merely an application of existing rectified flow techniques but is a novel approach, as we explain in the general response. While we acknowledge the potential of consistency models and similar diffusion-like techniques in imitation learning settings, our method stands out due to its efficient training and inference capabilities, as well as its adeptness at learning a spectrum of behaviors—features that are critical for practical, real-world applications. 

Furthermore, consistency models require distillation approaches to achieve low-NFE performance. RF-POLICY, on the other hand, achieves this efficiency intrinsically, without the need for such additional distillation or reflow steps. Future work may explore the possibility of using consistency models, but our current focus is to highlight the unique contributions of RF-POLICY and its practical implications.

Reviewer comment (Theory):
There is a theoretical gap between the objective at eq.~(8) and the implementation at Alg.1. The implementation uses a rectified flow to train the policy function, and uses another neural network to train the variance prediction network. In execution, the variance prediction network is used to determine the update iteration.

Author response:
**Q1. Gap between the objective at Eq.(8) and the implementation at Alg.1**

The two-stage training framework is indeed deliberate and advantageous. 

Firstly, training the variance prediction network in fact takes a very short time compared to the training of the RF network. Training the rectified flow model requires a duration that is 10 times longer than that needed to train the variance prediction model. 

Secondly, we empirically found that training the two components separately leads to slightly better performance than training them jointly (See the Table here ([link](https://anonymous.4open.science/r/RF-POLICY-184B/table4_separate_joint_train.png)), or Table 4 in the paper). 

|                    | Maze 1      |            |       |      | Maze 2      |            |       |      | Maze 3      |            |       |      | Maze 4      |            |       |      | Average     |            |       |      |
|--------------------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|
|                    | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  |
| RF-POLICY (Separate) | 1.00        | 0.98       | 1.11  | 0.91 | 1.00        | 0.79       | 1.03  | 0.83 | 1.00        | 0.96       | 1.99  | 0.78 | 0.98        | 0.93       | 1.85  | 0.59 | 0.99        | 0.91       | 1.50  | 0.78 |
| RF-POLICY (Joint)    | 1.00        | 0.97       | 1.09  | 0.67 | 1.00        | 0.93       | 1.08  | 0.76 | 1.00        | 0.91       | 2.08  | 0.64 | 1.00        | 0.91       | 1.70  | 0.60 | 1.00        | 0.93       | 1.49  | 0.67 |

Table 4: Performance comparison of separate training and joint training of RF-POLICY in Maze tasks. The table presents key performance metrics, including Success Rate (SR), Diversity Score (DS), Number of Function Evaluations (NFE), and Learning Efficiency Index (LEI), across various maze complexities. Results are shown for each model in different maze scenarios, as well as the average performance.

Therefore, though in theory training the two networks jointly or sequentially should be the same, in practice, training two “tasks” simultaneously can lead to optimization challenges (See PCGrad for the conflicting gradient issue when more objectives are introduced). Moreover, as we found that once the RF net is trained, training the variance network adds negligble extra overhead, so we choose this way of training. We have modified the paper to illustrate the reason in detail.

Reviewer comment (Experiments):
Considering that the variance prediction network and the policy are trained separately, the performance gain especially in training is only a contribution of the rectified flow instead of the proposed solution as a whole.

Author response:
**Q2:  Performance gain in training:**

First, regarding the performance in terms of the success rate and diversity, the benefit indeed comes from the RF network. It is in the original theory of RF that, if we ignore the simulation speed (how many steps to simulate the ODE), RF without any reflow performs the best. Applying RF on robotics problems is similar to applying DDIM/DDPM in the diffusion-policy paper. Moreover, we believe that RF is simpler to implement and easier to undertand, and is more computationally efficient.

However, we would like to argue that regarding the performance in terms of inference speed, the contribution mainly comes from the **variance network**. This variance network helps reduce the ODE simulation steps whenever possible. For instance, if a state always leads to a deterministic action, RF-policy will be as efficient as BC. And the number of simulation steps completely depends on the “variance” (or diversity) of the optimal actions given the state. This part alone is a core contribution.

Reviewer comment (Reproducibility):
The paper is clearly written with good visualizations. However, the gap between the objective at eq.(8) and the implementation is not explained.

Author response:
**Q1. Gap between the objective at Eq.(8) and the implementation at Alg.1**

The two-stage training framework is indeed deliberate and advantageous. 

Firstly, training the variance prediction network in fact takes a very short time compared to the training of the RF network. Training the rectified flow model requires a duration that is 10 times longer than that needed to train the variance prediction model. 

Secondly, we empirically found that training the two components separately leads to slightly better performance than training them jointly (See the Table here ([link](https://anonymous.4open.science/r/RF-POLICY-184B/table4_separate_joint_train.png)), or Table 4 in the paper). 

|                    | Maze 1      |            |       |      | Maze 2      |            |       |      | Maze 3      |            |       |      | Maze 4      |            |       |      | Average     |            |       |      |
|--------------------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|
|                    | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  |
| RF-POLICY (Separate) | 1.00        | 0.98       | 1.11  | 0.91 | 1.00        | 0.79       | 1.03  | 0.83 | 1.00        | 0.96       | 1.99  | 0.78 | 0.98        | 0.93       | 1.85  | 0.59 | 0.99        | 0.91       | 1.50  | 0.78 |
| RF-POLICY (Joint)    | 1.00        | 0.97       | 1.09  | 0.67 | 1.00        | 0.93       | 1.08  | 0.76 | 1.00        | 0.91       | 2.08  | 0.64 | 1.00        | 0.91       | 1.70  | 0.60 | 1.00        | 0.93       | 1.49  | 0.67 |

Table 4: Performance comparison of separate training and joint training of RF-POLICY in Maze tasks. The table presents key performance metrics, including Success Rate (SR), Diversity Score (DS), Number of Function Evaluations (NFE), and Learning Efficiency Index (LEI), across various maze complexities. Results are shown for each model in different maze scenarios, as well as the average performance.

Therefore, though in theory training the two networks jointly or sequentially should be the same, in practice, training two “tasks” simultaneously can lead to optimization challenges (See PCGrad for the conflicting gradient issue when more objectives are introduced). Moreover, as we found that once the RF net is trained, training the variance network adds negligble extra overhead, so we choose this way of training. We have modified the paper to illustrate the reason in detail.

Reviewer comment (Experiments):
I was wondering if there are any reasons supporting the implementation.

Author response:
**Q1. Gap between the objective at Eq.(8) and the implementation at Alg.1**

The two-stage training framework is indeed deliberate and advantageous. 

Firstly, training the variance prediction network in fact takes a very short time compared to the training of the RF network. Training the rectified flow model requires a duration that is 10 times longer than that needed to train the variance prediction model. 

Secondly, we empirically found that training the two components separately leads to slightly better performance than training them jointly (See the Table here ([link](https://anonymous.4open.science/r/RF-POLICY-184B/table4_separate_joint_train.png)), or Table 4 in the paper). 

|                    | Maze 1      |            |       |      | Maze 2      |            |       |      | Maze 3      |            |       |      | Maze 4      |            |       |      | Average     |            |       |      |
|--------------------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|-------------|------------|-------|------|
|                    | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  | SR          | DS         | NFE   | LEI  |
| RF-POLICY (Separate) | 1.00        | 0.98       | 1.11  | 0.91 | 1.00        | 0.79       | 1.03  | 0.83 | 1.00        | 0.96       | 1.99  | 0.78 | 0.98        | 0.93       | 1.85  | 0.59 | 0.99        | 0.91       | 1.50  | 0.78 |
| RF-POLICY (Joint)    | 1.00        | 0.97       | 1.09  | 0.67 | 1.00        | 0.93       | 1.08  | 0.76 | 1.00        | 0.91       | 2.08  | 0.64 | 1.00        | 0.91       | 1.70  | 0.60 | 1.00        | 0.93       | 1.49  | 0.67 |

Table 4: Performance comparison of separate training and joint training of RF-POLICY in Maze tasks. The table presents key performance metrics, including Success Rate (SR), Diversity Score (DS), Number of Function Evaluations (NFE), and Learning Efficiency Index (LEI), across various maze complexities. Results are shown for each model in different maze scenarios, as well as the average performance.

Therefore, though in theory training the two networks jointly or sequentially should be the same, in practice, training two “tasks” simultaneously can lead to optimization challenges (See PCGrad for the conflicting gradient issue when more objectives are introduced). Moreover, as we found that once the RF net is trained, training the variance network adds negligble extra overhead, so we choose this way of training. We have modified the paper to illustrate the reason in detail.

Reviewer comment (Experiments):
I was wondering if using rectified flow to replace the DDPM in the DDIM models will lead to similar performances in both tasks.

Author response:
**Q3. Replacing DDPM with rectified flow in DDIM models:**

In response to your query, we have implemented the standard rectified flow method as an alternative to DDIM in our experiments. Our comparison includes both the 1-Rectified Flow (1-RF, before reflow) and the 2-Rectified Flow (2-RF, after reflow), using a standard Euler sampler for action generation. (See the Table here ([link](https://anonymous.4open.science/r/RF-POLICY-184B/table5_compare_with_rf.png)), or Table 5 in the paper). 

Our findings indicate that while 1-RF shows limitations in capturing diverse behaviors in a single step, 2-RF demonstrates an ability to produce varied actions with just one step. However, it's important to note that 2-RF involves a reflow process which is approximately 7 times more resource-intensive in terms of training compared to 1-RF, Behavioral Cloning (BC), or our RF-POLICY method.

|                  | Maze 1 |      |      |      | Maze 2 |      |      |      | Maze 3 |      |      |      | Maze 4 |      |      |      | Average |      |      |      |
|------------------|--------|------|------|------|--------|------|------|------|--------|------|------|------|--------|------|------|------|---------|------|------|------|
|                  | SR     | DS   | NFE  | LEI  | SR     | DS   | NFE  | LEI  | SR     | DS   | NFE  | LEI  | SR     | DS   | NFE  | LEI  | SR      | DS   | NFE  | LEI  |
| 1-RF (step=1)    | 1.00   | 0.95 | 1.00 | 0.94 | 1.00   | 0.93 | 1.00 | 0.92 | 1.00   | 0.26 | 1.00 | 0.84 | 0.96   | 0.59 | 1.00 | 0.75 | 0.99    | 0.68 | 1.00 | 0.86 |
| 1-RF (step=5)    | 1.00   | 1.00 | 5.00 | 0.93 | 0.92   | 1.00 | 5.00 | 0.75 | 1.00   | 0.99 | 5.00 | 0.83 | 0.92   | 1.00 | 5.00 | 0.84 | 0.96    | 1.00 | 5.00 | 0.84 |
| 1-RF (step=20)   | 1.00   | 1.00 | 20.00| 0.92 | 0.82   | 1.00 | 20.00| 0.68 | 1.00   | 1.00 | 20.00| 0.84 | 0.94   | 1.00 | 20.00| 0.86 | 0.94    | 1.00 | 20.00| 0.83 |
| 2-RF (step=1)    | 1.00   | 1.00 | 1.00 | 0.12 | 0.86   | 0.98 | 1.00 | 0.09 | 1.00   | 1.00 | 1.00 | 0.12 | 0.98   | 1.00 | 1.00 | 0.11 | 0.96    | 1.00 | 1.00 | 0.11 |
| 2-RF (step=5)    | 1.00   | 1.00 | 5.00 | 0.12 | 0.86   | 1.00 | 5.00 | 0.10 | 1.00   | 1.00 | 5.00 | 0.12 | 0.98   | 1.00 | 5.00 | 0.11 | 0.96    | 1.00 | 5.00 | 0.11 |
| 2-RF (step=20)   | 1.00   | 1.00 | 20.00| 0.12 | 0.88   | 1.00 | 20.00| 0.10 | 1.00   | 1.00 | 20.00| 0.12 | 0.96   | 1.00 | 20.00| 0.11 | 0.96    | 1.00 | 20.00| 0.11 |
| RF-POLICY        | 1.00   | 0.98 | 1.11 | 0.91 | 1.00   | 0.79 | 1.03 | 0.83 | 1.00   | 0.96 | 1.99 | 0.78 | 0.98   | 0.93 | 1.85 | 0.59 | 0.99    | 0.91 | 1.50 | 0.78 |