Normalizing Flows are Capable Models for Bi-manual Visuomotor Policy

1. Problem Statement

The field of general-purpose robotics has increasingly adopted Diffusion Models for learning visuomotor policies (mapping visual observations to robot actions). While diffusion models offer high expressiveness for multi-modal distributions, they suffer from two critical limitations:

1. High Computational Cost: Inference requires expensive, multi-step denoising processes, making real-time control on resource-constrained hardware difficult.
2. Lack of Uncertainty Quantification: Diffusion models do not provide tractable estimates of output likelihood, preventing the system from quantifying confidence or optimizing trajectories based on probability density.

Existing regression-based models (e.g., Transformers) are fast but often fail to capture multi-modal distributions, collapsing to the mean of conflicting demonstrations. The authors propose Normalizing Flows (NFs) as an alternative that combines the expressiveness of generative models with the efficiency of single-pass sampling and exact likelihood computation.

2. Methodology: Normalizing Flow Policy (NF-P)

The authors introduce NF-P, a conditional normalizing flow architecture designed for bi-manual manipulation.

Core Architecture

• Bijective Mapping: NF-P learns a bijective, differentiable transformation $T$ that maps complex action sequences ($x$) to a simple Gaussian prior ($z$), conditioned on visual observations ($o$).
• Likelihood Computation: Due to the change-of-variables formula, the model can compute the exact log-likelihood $\log p(x|o)$ in a single forward pass.
• Encoding: Visual inputs are processed via a pre-trained ResNet18. The resulting embedding, combined with the latest action, conditions the flow network.
• Flow Structure: The model utilizes Neural Spline Flows (rational-quadratic splines) within coupling layers to ensure expressiveness and invertibility.

Key Innovations

1. Strided Sampling (Training):
   • Problem: Human teleoperation data contains high-frequency noise (jitters, pauses) that distracts the model.
   • Solution: Instead of standard sliding windows, the authors apply a temporal stride $s$ to both observation and action sequences during training. This forces the model to learn transitions between larger state changes, effectively filtering out unconscious noise while maintaining dense dataset coverage.
2. Inference-Time Optimization Strategies:
   Leveraging the exact likelihood capability, the authors propose two strategies to refine action selection:
   • Stochastic Batch Selection (SBS): Generates a batch of $N$ candidate trajectories (e.g., $N=128$) and selects the one with the highest log-likelihood. This filters out low-probability outliers.
   • Gradient Refinement (GR): Takes an initial sample and performs gradient ascent on the log-likelihood ($\nabla_x \log p(x|o)$) for a few steps (e.g., 10 iterations) to "climb" the probability manifold toward a local maximum, refining the trajectory without discarding the stochastic nature of the policy.

3. Key Contributions

1. Novel Architecture: Proposal of NF-P, a conditional normalizing flow tailored for RGB-based bi-manual policies, featuring strided data sampling and likelihood-based output optimization.
2. Performance Superiority: Demonstration that NF-P matches or exceeds the success rates of Diffusion Policy (DP) baselines while offering significantly faster training and inference.
3. Scalability: Evidence that NFs scale effectively with dataset size, comparable to diffusion models.
4. Real-World Robustness: Validation on a physical dual-arm Kuka setup, showing NF-P's ability to overcome transition bottlenecks and stalling issues common in diffusion-based approaches.

4. Experimental Results

Simulation (RoboTwin 2.0)

• Benchmark: Evaluated on 50 bi-manual tasks (grasping, tool use, stacking).
• Success Rate: NF-P with Gradient Refinement (NF-PGR) achieved an average success rate of 38.06%, a 10% absolute improvement over the Diffusion Policy baseline (28.04%).
• Precision Tasks: The gap widened in high-precision tasks (e.g., Open Microwave: 77% vs. 5%; Handover Mic: 100% vs. 53%).
• Efficiency: NF-P trained in 50 minutes (300 epochs) on 100 demonstrations. Inference time was <20ms for SBS and **455ms** for GR (still faster than standard DP).

Real-World Robot (Dual-Arm Kuka)

• Tasks: Stack Block Two and Towel Folding (deformable object).
• Data: 150 and 100 demonstrations respectively.
• Bottleneck Analysis:
  • Diffusion Policy: Suffered from high initial failure rates (35% for stacking) and frequently got stuck in sub-goal transitions (e.g., failing to switch arms).
  • NF-P: Achieved 0% initial failure for stacking and significantly higher overall success rates (60% vs. 15% for stacking; 75% vs. 55% for towel folding).
  • Transition Speed: NF-P completed sub-goal transitions in ~30 seconds, whereas DP required up to 8 minutes before failing.

Ablation Studies

• Stride ($s$): $s=4$ was optimal. $s=1$ (no stride) led to noise sensitivity (37% success), while $s=8$ was too aggressive (46%).
• Sampling Variance ($\sigma$): Moderate variance ($\sigma=0.5$) yielded the best results, balancing exploration and stability.
• Optimization: Both SBS and GR significantly outperformed direct sampling.

5. Significance and Conclusion

This work establishes Normalizing Flows as a highly competitive alternative to diffusion models for robotic policy learning. The significance lies in three main areas:

1. Real-Time Viability: By enabling single-pass inference, NF-P removes the latency barrier of diffusion models, making them suitable for high-frequency control loops.
2. Uncertainty-Aware Control: The ability to compute exact likelihoods allows for principled uncertainty estimation and trajectory optimization (via gradient ascent), leading to safer and more robust exploration.
3. Data Efficiency: NF-P demonstrates strong performance with fewer training epochs and data compared to diffusion baselines, particularly in complex, multi-modal bi-manual tasks.

The authors conclude that NF-P provides a mathematically sound, computationally efficient framework for scaling robotic learning to real-world, multi-modal settings where responsiveness and reliability are paramount.