Technical Summary: Reverse Flow Matching (RFM)

1. Problem Statement

The paper addresses a fundamental challenge in applying Diffusion and Flow policies to Online Reinforcement Learning (RL). While these generative models have achieved success in offline RL and imitation learning due to their expressive power, their deployment in online RL is hindered by the lack of direct samples from the target policy distribution.

In the Maximum Entropy RL framework, the optimal policy is defined by a Boltzmann distribution over actions:
$$\pi_{\text{new}}(a \mid s) \propto \exp\left(\frac{1}{\lambda} Q(s, a)\right)$$
Unlike standard generative modeling where training data is readily available, this target distribution is unnormalized and generally intractable to sample from directly. Existing approaches to circumvent this have fallen into two distinct families for diffusion policies:

1. Noise-Expectation: Uses a weighted average of noise samples (e.g., Q-weighted noise estimation).
2. Gradient-Expectation: Uses a weighted average of Q-function gradients (e.g., Iterated Denoising Energy Matching).

The paper identifies that the formal relationship between these two objectives remains unclear, and it is unknown if they can be synthesized into a more general formulation. Furthermore, existing methods are often rigidly coupled to specific noise schedules and have been limited to diffusion policies, leaving the training of Flow policies to sample from Boltzmann distributions as an open problem.

2. Methodology: Reverse Flow Matching (RFM)

The authors propose Reverse Flow Matching (RFM), a unified framework that reframes the training of diffusion and flow models without direct target samples as a posterior mean estimation problem.

2.1. From Forward Construction to Reverse Inference

Standard conditional flow matching relies on a forward construction: sampling a source $X_0 \sim p_0$ and a target $X_1 \sim p_1$, then interpolating to $X_t$. RFM reverses this perspective. Since $X_1$ (the action from the Boltzmann distribution) cannot be sampled directly, the framework treats the intermediate state $X_t$ as observed evidence and infers the latent variables ($X_0$ or $X_1$) that generated it.

Given an interpolation $X_t = \alpha_t X_1 + \beta_t X_0$, the posterior distribution of the noise $X_0$ given $X_t$ is:
$$q^*_{0|t}(x_0 \mid x_t) \propto p_0(x_0) p_1\left(\frac{1}{\alpha_t}x_t - \frac{\beta_t}{\alpha_t}x_0\right)$$
The training objective is transformed into regressing the velocity field $v_\theta$ toward the posterior mean of the velocity, which is a function of the posterior mean of $X_0$ or $X_1$.

2.2. Langevin Stein Operators and Control Variates

The core computational challenge is estimating the posterior mean $\mathbb{E}[X_0 \mid X_t]$ when the target density $p_1$ is known only up to a normalizing constant. The paper introduces Langevin Stein operators to construct zero-mean control variates that reduce the variance of these estimators.

• Control Variate Construction: By leveraging the property that the expectation of a Langevin Stein operator applied to a sufficiently regular function is zero, the authors construct vector-valued control variates $g_{\Phi_t}$.
• Variance Reduction: The posterior mean estimator is modified to:
  $$\hat{\mu} = \mathbb{E}_{X_0 \sim \bar{q}} \left[ \frac{w(X_0) (X_0 + g_{\Phi_t}(X_0))}{\sum w(X_0)} \right]$$
  where $w$ are importance weights and $g_{\Phi_t}$ is the control variate derived from the Stein operator.
• Unified Estimators: The authors derive a general class of estimators parameterized by a test function $\Phi$.
  • Setting the parameter $\eta = 0$ recovers the Noise-Expectation family.
  • Setting $\eta = 1$ recovers the Gradient-Expectation family.
  • The framework allows for the principled combination of Q-values and Q-gradients to form a variance-reduced estimator.

2.3. Application to Flow Policies in Online RL

The framework is instantiated for Flow policies in continuous-control tasks:

• Latent Space Mapping: To handle action bounds (e.g., $a \in [-1, 1]^d$), the flow is learned in an unconstrained latent space $u$ where $a = \tanh(u)$. This avoids heuristic truncation that breaks probability paths.
• Jacobian Correction: The target density in the latent space includes a Jacobian factor from the change of variables, ensuring the correct Boltzmann distribution in action space.
• Training Loop: The actor is trained by minimizing the squared error between the predicted velocity and the RFM target (the posterior mean velocity estimated via Self-Normalized Importance Sampling with control variates).

3. Key Contributions

1. Tractable RFM Loss: The paper proposes a loss function that transforms the intractable problem of unavailable target samples into a posterior mean estimation problem, applicable to both diffusion and flow models.
2. General Estimator Class via Stein Operators: By introducing Langevin Stein operators, the authors derive a general class of posterior mean estimators. They demonstrate that existing noise-expectation and gradient-expectation methods are merely specific instances ($\eta=0$ and $\eta=1$) of this broader class. This enables the combination of Q-value and Q-gradient information to improve training stability.
3. Extension to Flow Policies: The framework extends the capability of targeting Boltzmann distributions from diffusion models to Flow policies. This leverages the flexibility of flow models to admit general source distributions beyond standard Gaussian noise.
4. Empirical Validation: The authors instantiate RFM for online RL and demonstrate superior performance and stability compared to diffusion policy baselines on continuous-control benchmarks.

4. Experimental Results

The paper evaluates RFM on two fronts:

• Toy Example (2D Two-Moon Distribution): RFM was compared against diffusion baselines (iDEM and QNE). RFM achieved the lowest distribution discrepancy (Sliced Wasserstein, MMD, Sinkhorn) while requiring only 20 inference steps compared to 100 steps for the baselines, despite using the same training computation budget.
• Continuous-Control Benchmarks (DeepMind Control Suite): RFM was tested on eight environments against five baselines: SAC, Q-score matching (QSM), MaxEntDP (QNE), Diffusion Q-sampling (DQS), and QVPO.
  • Stability: RFM was the only algorithm to perform consistently well across all eight environments. Baselines struggled significantly on specific tasks.
  • Efficiency: RFM used only 10 flow steps for inference, whereas diffusion baselines required 20 steps.
  • Performance: RFM achieved higher total rewards and greater stability than all baselines.

5. Significance and Claims

The paper claims that Reverse Flow Matching provides a unified theoretical framework that rigorously connects previously distinct approaches to training generative policies in online RL. By adopting a reverse inferential perspective, it resolves the issue of training without direct target samples.

The significance lies in:

• Unification: It clarifies that noise-expectation and gradient-expectation methods are not competing paradigms but specific cases of a general posterior mean estimation problem.
• Stability and Efficiency: The use of Langevin Stein operators for variance reduction leads to more stable training and allows for fewer inference steps (flow steps) compared to diffusion baselines.
• Generalization: It successfully extends these capabilities to Flow policies, opening the door for incorporating domain knowledge through tailored source distributions and improving the practical deployment of expressive policies in complex continuous-control tasks.

The authors maintain a modest tone, noting that while the framework unifies existing methods, the specific choice of proposal distributions and control variate parameters can be further optimized, and the method is currently instantiated for online RL with off-policy critics.