RF-POLICY: Rectified Flows are Adaptive Decision Makers

We sincerely thank the reviewers for their insightful comments and constructive suggestions, which have been instrumental in enhancing the clarity and depth of our work.

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

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.

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.

We sincerely thank the reviewers for their insightful comments and constructive suggestions, which have been instrumental in enhancing the clarity and depth of our work.

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) 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: Theorem 1 proof:

Theorem 1 proves that for deterministic behaviors, where the variance of action given state is zero, the RF-POLICY loss function is optimized such that the flow model will produce straight-line trajectories. However, this does not limit RF-POLICY's ability to generate multi-modal behaviors. There’s an expectation over the optimal velocity, $v^*(z, t | x)$, as outlined in the proof, captures the expected trajectory.

In the multi-modal case, this expectation accounts for the multiple possible directions $y - Z_0$ that can pass through point $z$ at time $t$. Consequently, for a multi-modal distribution, the variance would be greater than zero, allowing for the generation of diverse trajectories.

Our empirical evaluations support the model's multi-modal generative capabilities. In scenarios with inherent uncertainty and multiple valid actions for a given state, RF-POLICY successfully generates a diverse set of plausible actions, which is proof of its ability to model multi-modal behavior effectively.

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.

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

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 ​

We sincerely thank the reviewers for their insightful comments and constructive suggestions, which have been instrumental in enhancing the clarity and depth of our work.

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 or Figure 7 in our paper). The table shows the comparison between bc, standard RF, Rectified Flow and RF-POLICY (See this link or Table 8 in our paper).

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.

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 or Table 7 in our paper).

We sincerely thank the reviewers for their insightful comments and constructive suggestions, which have been instrumental in enhancing the clarity and depth of our work.

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), or Table 4 in the paper).

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.

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.

Theory:

1. In the original manuscript, we have established the result that RF-POLICY yields a straight-line ODE process if $P(a | s)$ is deterministic.
2. In the updated method, we present a further result showcasing that the number of simulation steps is determined by how “multimodal” or “diverse” $P(a | s)$ is. And the resulting simulation error is strictly bounded. Note that 2 is a compliment result of 1 and not contradicting 1.

Comparison to Offline RL diffusion-based methods: Reviewer J2ZQ and KiF7 were curious about how RF-POLICY compares against offline RL methods like the Diffuser, Decision Diffuser, CQL, Decision Diffuser, etc. Here we note that in principle, we can definitely extend RF-policy’s idea to the offline RL setting, but RF-POLICY, just like Diffusion Policy, is a discriminative model that only predicts $P(a | s)$. There is no assumption about if offline rewards are provided, nor do we try to predict the joint (s, a) distribution. Hence, RF-POLICY in principle still remains to be an offline imitation learning method (offline as opposed to online method like GAIL/ AIRL that needs interaction with the environment).

Empirical Validation: In response to reviewers' feedback, we have conducted additional experiments, which encompass the following aspects:

1. Comparison with Standard RF Method: We have performed a thorough comparison with the standard Random Forest (RF) method, considering both 1-RF (before reflow) and 2-RF (after reflow) models.
2. Ablation Analysis with Low-NFE: We have extended our ablation study to include results using low Neural Function Evaluation (NFE) for both DDIM and the standard RF method.
3. Efficiency Metrics for Training and Inference: We have introduced metrics to assess the efficiency of our models during both training and inference stages.
4. Empirical Validation of Assumption 1: We have provided empirical validation for Assumption 1 by visualizing the variance in actions observed in robot demonstrations.
5. Multi-Modal Behavior in Robot Manipulation: We have gathered demonstrations for a pick-and-place robot manipulation task, where the robot employs various poses to pick up objects. We subsequently compared different methods in this multi-modal context.
6. Comparison with Other Sampling Methods: In the domain of image generation, we have conducted a comparison between RF-POLICY and alternative sampling methods.

We have reflected the extended experiments in the updated version of the submission. Please see the appendix for extended experiments.