Drift Q-Learning

Technical Summary: Drift Q-Learning (DriftQL)

Problem Statement
Offline Reinforcement Learning (RL) aims to learn optimal policies from fixed datasets without further environment interaction. A central challenge in this setting is the "out-of-distribution" (OOD) problem: value functions trained on static data produce unreliable estimates for actions not present in the dataset. To mitigate this, behavior regularization is required to constrain the learned policy toward the data distribution while still allowing for improvement toward high-value actions.

Existing solutions utilizing expressive generative models, such as diffusion and flow-based policies, effectively model the multimodal structure of behavior distributions. However, they suffer from significant practical limitations: they require iterative denoising or heavy solver integrations at inference time, leading to slow execution. While distillation-based alternatives exist to speed up inference, they introduce complexity by requiring additional networks and two-stage training pipelines.

Methodology
The authors propose Drift Q-Learning (DriftQL), a one-step generative method that combines a drift-based behavioral regularizer with critic-driven policy improvement. DriftQL is built upon the framework of "Drifting Models," which learn a direct one-step pushforward map from prior noise and conditioning signals to data.

The core of DriftQL is a stochastic actor network $f_\theta(s, \epsilon)$ that maps a state $s$ and Gaussian noise $\epsilon$ to an action $a$. The training objective unifies distribution modeling with reward seeking through two primary mechanisms:

1. Conditional Drift Field: The policy is supervised by a state-conditional drift field composed of two opposing forces:

   • Attraction ($V^+$): Pulls generated actions toward the observed dataset actions (behavior cloning). In the offline RL setting, where there is typically one observed action per state, this simplifies to a direct vector difference between the generated action and the dataset action.
   • Repulsion ($V^-$): Pushes generated actions away from one another. This prevents the policy from collapsing onto a single mode (the dataset action) and maintains a diverse set of candidates within high-value regions. The repulsion is calculated using a kernel-weighted mean-shift over the $N$ candidate actions generated for a single state.

2. Critic-Driven Improvement: Unlike standard drifting models used for unconditional generation, DriftQL incorporates a value-based objective. A clipped double-Q critic estimates the value of actions. The actor minimizes a combined loss:
   $$L_{actor}(\theta) = \alpha L_{drift}(\theta) + L_Q(\theta)$$
   where $L_{drift}$ is the drift loss (regressing actions toward the drifted target) and $L_Q$ is the negative expected Q-value. The value gradient biases the policy toward high-value regions of the data support, allowing the policy to concentrate on multiple high-value modes if they exist in the data.

Key Contributions

• One-Step Generative Policy: DriftQL generates actions in a single forward pass, eliminating the need for iterative denoising chains, ODE solvers, or distillation stages required by diffusion and flow-based methods.
• Unified Training Objective: The method integrates behavior regularization (via the drift field) and policy improvement (via the Q-critic) into a single network with a unified loss, avoiding the complexity of multi-stage training pipelines.
• Mode Selection and Diversity: The combination of attraction (to data) and repulsion (between candidates) allows the policy to maintain mass across candidates during training while converging to high-value modes, preventing collapse to a single deterministic action or scattering across unsupported regions.

Experimental Results
DriftQL was evaluated on the D4RL and OGBench benchmarks, covering continuous control, navigation, and manipulation tasks.

• Performance: DriftQL consistently outperforms or matches state-of-the-art diffusion and flow-based baselines (e.g., Diffusion-QL, FQL, IDQL) across 78 offline learning tasks. It shows particularly strong gains on long-horizon navigation tasks (e.g., antmaze-large-st, antmaze-giant-st) and complex manipulation tasks.
• Robustness: Under degraded data quality (simulated by replacing a fraction of actions with random noise), DriftQL maintains performance closer to its clean-data baseline than competing methods, which visibly struggle as noise increases.
• Efficiency: In terms of inference latency, DriftQL is approximately 2$\times$ faster than FQL, 3$\times$ faster than Diffusion-QL, and 4$\times$ faster than IDQL/IFQL. It achieves this without requiring a separate distillation stage. Training costs scale with the number of generated samples ($N$) due to pairwise kernel calculations but remain modest compared to the overhead of iterative denoising in diffusion models.

Significance and Claims
The paper positions DriftQL as a promising alternative to diffusion and flow-based offline RL methods. It claims to offer the expressivity of generative models (handling multimodal distributions) while retaining the simplicity and efficiency of deterministic approaches (single forward pass).

The authors argue that DriftQL successfully reconciles the competing objectives of offline RL: it constrains the policy to the data support to ensure reliable value estimates while permitting improvement toward high-reward regions. By removing the computational overhead of iterative sampling, DriftQL provides a more practical solution for real-world deployment where inference speed is critical, all while maintaining robustness against data corruption. The work suggests that drifting provides a useful "middle ground" for offline RL, offering stochastic, distributional training signals without the test-time costs associated with multi-step generative policies.