Q-learning with Flow-Matching Policies

The Rise of Expressive Policies in Robotics

Recent progress in learning-based robotic manipulation, from single-task demonstrations to multi-task generalist systems, has been driven by several key ingredients. While large-scale datasets and pre-trained VLMs are crucial, two algorithmic components stand out:

1. Action Chunking: This technique involves grouping a sequence of control actions into a single “chunk,” which a policy predicts and executes together. This approach, motivated by neuroscience, offloads cognitive burden, improves temporal coherency, and better handles partial observability. Crucially, it helps model complex human demonstration data, which often violates the Markov assumption due to memory and temporal correlation.
2. Generative Policies (Diffusion & Flow-Matching): A challenge introduced by action chunking is that the distribution of action chunks becomes highly complex and multi-modal. Traditional Gaussian policies are insufficient for modeling these distributions. The adoption of diffusion and flow-matching policies has been a breakthrough, as they excel at modeling such multi-modal distributions and can generalize effectively from a small number of demonstrations.

From Imitation Learning to Reinforcement Learning

The standard recipe for training these policies is supervised imitation learning (IL) on large-scale prior experience. While effective, IL struggles to achieve the last mile of performance and reliability (e.g., going from 90% to 100% success). To achieve this, robots need to learn from “counterfactuals” and online experience, which are costly and difficult to collect purely through imitation.

This necessitates the use of Reinforcement Learning (RL) for online self-improvement and scaling data collection. The talk focuses on Q-learning, an off-policy RL framework well-suited for this problem because it can leverage large-scale prior (offline) experience alongside new online data.

A standard Q-learning algorithm iteratively performs two steps:

• Policy Evaluation: Estimating the Q-value function (how good an action is in a given state) using the Bellman backup.
• Policy Improvement: Updating the policy to maximize the estimated Q-value.

For continuous control, efficient Q-learning algorithms utilize the reparameterization gradient, allowing the policy to be optimized effectively using the gradient of the Q-function. This works well for simple policies like Gaussians but poses a significant challenge for multi-step generative models like flow-matching policies.

Challenge: Unstable Gradients in Flow-Policy Optimization

The core problem is that naively optimizing a flow policy by backpropagating the Q-function’s gradient through the entire multi-step denoising process is highly unstable. Much like training a deep recurrent neural network, this can lead to exploding or vanishing gradients, hindering effective policy improvement. The central question is: “How can we leverage the expressive multi-step flow model while not suffering from this unstable backpropagation issue?”

Solution 1: FQ-RL - Q-Learning via Distillation

The first proposed algorithm, Flow Q-Learning (FQ-RL), addresses the instability problem through distillation.

• Core Idea: Instead of backpropagating through the complex multi-step flow model, distill its mapping into a stable, single-step model.
• Mechanism:
  1. A multi-step flow policy is trained via imitation on prior data.
  2. This policy is used to generate paired data of (noise, action).
  3. A separate single-step neural network (the “actor”) is trained to approximate this direct noise-to-action mapping.
  4. The actor’s loss function then simultaneously maximizes the Q-value (via stable backpropagation through the single-step network) and minimizes the distance to the expressive distribution generated by the original multi-step flow model.

Results: FQ-RL performs exceptionally well, especially in action chunking settings. On challenging manipulation tasks, it was the only method to show significant online improvement, highlighting the synergy between an expressive policy class and an efficient RL optimizer.

Drawback: The primary limitation is the potential loss of expressivity. Distilling a multi-step model into a single-step one can sacrifice some of the generative power and generalization capabilities, even if the single-step model is made significantly larger.

Solution 2: CAM - Q-Learning with Conditional Adjoint Matching

The second algorithm improves upon FQ-RL by avoiding distillation altogether, using a technique called Adjoint Matching (CAM).

• Core Idea: Leverage a more stable gradient estimation technique from the adjoint matching literature to avoid unstable backpropagation while still using the full, multi-step flow model.
• Mechanism:
  • The optimization problem is framed as a KL-constrained objective to maximize the Q-value while staying close to a pre-trained behavior policy.
  • Standard backpropagation through the flow model’s ODE depends on the fine-tuned velocity field, which is constantly changing and can be ill-conditioned, causing instability.
  • Adjoint matching provides a “linear adjoint” objective. Its key property is that the backpropagated gradient depends on the pre-trained behavior velocity field, not the fine-tuned one.
  • Since the behavior model is trained on a static dataset and is generally well-conditioned, the gradient computation becomes stable.

Results: CAM consistently outperforms other methods, including FQ-RL, across a wide range of offline and online fine-tuning tasks. It demonstrates superior sample efficiency in the challenging action chunking setting, showing the benefits of using the full expressive power of the flow model with stable Q-gradients.

Drawback: While much more stable, CAM can still occasionally suffer from instability if the underlying pre-trained behavior velocity field is itself ill-conditioned.

Future Directions

1. Hybrid Optimization: Develop algorithms that can intelligently combine gradient-based optimization (when Q-gradients are reliable) with sampling-based methods (when gradients are poor). This would allow the agent to pick the best action by sampling when the Q-landscape is noisy or non-convex, while still exploiting efficient gradients when available.
2. Improving Sample Efficiency: A key challenge remains the sample efficiency of online learning. The current methods require hours of real-robot interaction to show meaningful progress. A major goal is to reduce this fine-tuning time to the order of minutes, which would dramatically accelerate real-world deployment and scaling.

• Tags:
• Imitation Learning
• Reinforcement Learning
• Feature
• Robot Learning
• Q-Learning
• Diffusion Policies
• Flow Matching
• Action Chunking