Moment Matching Q-Learning

Score-based and flow-based generative models exhibit remarkable expressive capacity in capturing complex distributions, and have been extensively deployed in tasks ranging from image generation to reinforcement learning. Nevertheless, these models suffer from prolonged inference latency, which imposes a significant computational bottleneck in RL with iterative sampling. To overcome this limitation, we propose a new framework named Moment Matching Q-Learning (MoMa QL), which utilizes a technique from statistical hypothesis testing known as maximum mean discrepancy (MMD) that intend to match all orders of statistics between the original and target distribution. By enforcing strong regularization on all moment statistics, this algorithm guarantees distribution-level convergence for conditional score function and remains stable under various hyperparameters. Empirically, we show that our method MoMa QL is more computationally efficient with a comparable if not competitive performance in various D4RL tasks. Remarkably, by accelerating the action sampling process for flow-based policies, MoMa QL demonstrates superior performance in offline-to-online RL tasks because of faster and stronger adaptability for online interactive finetuning.

Robot agents have been proved to have ability to learn complex skills by watching pre-recorded data instead of risky trial-and-error in the real world. While advanced generative techniques help robots learn highly sophisticated behaviors, their deployment remains constrained by computational latency; generating a single action requires dozens of repetitive computational steps, making them impractical for real-time systems like self-driving cars or agile manufacturing. To break this speed bottleneck, we introduce Moment Matching Q-Learning (MoMa QL). Our framework uses a mathematical alignment technique called Maximum Mean Discrepancy (MMD) to match the robot's decisions with expert data across all statistical profiles simultaneously. This enables the AI to capture complex, flexible behaviors safely and generate high-quality decisions in just 2 to 4 rapid steps, preventing policy collapse without relying on slow iterative processes. Experimentally, MoMa QL trains and runs up to 6 times faster than traditional generative methods while matching or exceeding the top-tier performance of current state-of-the-art algorithms across major robotic benchmarks. This breakthrough paves the way for a future of highly responsive, computationally efficient, and safe autonomous systems.