Fisher Decorator: Refining Flow Policy via A Local Transport Map

1. Problem Statement

Offline Reinforcement Learning (RL) aims to learn optimal policies from fixed datasets without further environment interaction. A central challenge is balancing reward maximization with distributional shift mitigation. The standard approach formulates this as a KL-constrained optimization problem, ensuring the learned policy ($\pi_\theta$) stays close to the behavioral policy ($\pi_\beta$) defined by the dataset.

Recent advances utilize Flow-based models (via Flow Matching) to parameterize policies, offering high expressiveness for complex, multimodal distributions. However, existing flow-based offline RL methods (e.g., Flow Q-Learning, DeFlow) face a critical theoretical and practical trade-off:

• The Mismatch: These methods typically approximate the KL constraint using an $L_2$ regularization term, which serves as an upper bound for the 2-Wasserstein distance ($W_2$).
• The Flaw: The KL divergence is inherently anisotropic (direction-dependent) and density-aware; it strictly penalizes deviations into low-density regions (support preservation). In contrast, the $L_2$/$W_2$ approximation is isotropic (direction-independent) and density-insensitive.
• Consequences: This geometric mismatch leads to:
  1. Mode Collapse/Averaging: The policy collapses toward the geometric mean of modes rather than preserving multimodality.
  2. Distributional Shift: The policy drifts into low-density or unsupported regions (Out-of-Distribution) because the isotropic penalty fails to distinguish between high-density and low-density deviations.
  3. Suboptimal Updates: Optimization directions are misaligned with the intrinsic geometry of the data manifold.

2. Methodology: Fisher Decorator (FiDec)

The authors propose Fisher Decorator (FiDec), a framework that reframes policy refinement from a geometric perspective to resolve the KL-$W_2$ mismatch.

A. Local Transport Map Formulation

Instead of learning a completely new one-step flow policy, FiDec models the refined policy $\pi_\theta$ as a local transport map applied to the behavioral flow policy $\pi_\beta$.

• Definition: $\pi_\theta(\cdot|s) = (T_s)_\# \pi_\beta(\cdot|s)$, where $T_s(a) = a + \delta(s, a)$.
• Here, $\delta(s, a)$ is a learned residual displacement field. This ensures the refined policy is anchored to the support of the behavioral policy, naturally preventing mass from leaking into "voids."

B. Geometric Approximation of KL Divergence

By analyzing the pushforward measure induced by the transport map, the authors derive a second-order local quadratic approximation of the KL divergence:
$$D_{KL}(\pi_\theta \| \pi_\beta) \approx \frac{1}{2} \mathbb{E}_{a \sim \pi_\beta} \left[ \delta(s, a)^\top I(s, a) \delta(s, a) \right]$$

• Key Insight: The constraint is governed by the Local Fisher Information Matrix (FIM), $I(s, a) = \nabla_a \log \pi_\beta(a|s) \nabla_a \log \pi_\beta(a|s)^\top$.
• Significance: Unlike the isotropic $L_2$ norm, the FIM provides an anisotropic metric. It penalizes displacements more heavily in directions where the log-density gradient is steep (high sensitivity) and allows more freedom in flat directions. This aligns optimization with the statistical manifold of the data.

C. Tractable Estimation via Flow Velocity

A major hurdle is computing the score function $\nabla_a \log \pi_\beta$ required for the FIM. The authors leverage the Flow Matching framework:

• The score function is intrinsically linked to the learned velocity field $v_\beta(t, s, a)$.
• They derive a numerically stable estimator using a perturbed time step $t_\epsilon = 1 - \epsilon$:
  $$\nabla_a \log \pi_\beta(a|s) \approx \frac{t_\epsilon v_\beta(t_\epsilon, s, a) - a}{1 - t_\epsilon}$$
• This allows the FIM to be estimated efficiently without backpropagation through time or solving ODEs.

D. Optimization Algorithm

FiDec solves a trust-region optimization problem using a primal-dual approach:

1. Primal Update: Maximize the expected Q-value subject to the Fisher-based quadratic constraint.
2. Dual Update: Update the Lagrange multiplier $\lambda$ to enforce the constraint threshold $\epsilon$.
3. Efficiency: The method avoids iterative distillation or ODE solving, performing refinement via a simple one-step residual update.

3. Key Contributions

1. Geometric Characterization of the KL-$W_2$ Gap: The paper identifies that the failure of prior flow-based methods stems from approximating an anisotropic KL constraint with an isotropic $W_2$ bound. It proves that policy refinement naturally evolves on a statistical manifold governed by the Fisher metric.
2. Tractable Anisotropic Optimization: The authors reformulate the problem into a quadratic objective defined by the Fisher Information Matrix. They show this metric can be efficiently estimated from the flow velocity field, enabling metric-aware updates that preserve multimodality and support.
3. The Fisher Decorator (FiDec) Framework: A practical algorithm that parametrizes refinement as a local transport map. It achieves state-of-the-art performance without the computational overhead of iterative generation or ODE solving.

4. Experimental Results

FiDec was evaluated on diverse benchmarks, including OGBench (Open X-Embodiment) and D4RL, covering navigation, locomotion, and manipulation tasks.

• Offline Performance: FiDec consistently outperforms state-of-the-art baselines (including Gaussian, Diffusion, and Flow-based methods like FQL and DeFlow).
  • Multimodal Tasks: In complex, multimodal environments (e.g., Humanoid, Puzzle, AntSoccer), FiDec successfully preserves multiple modes, whereas baselines suffer from mode collapse or averaging.
  • Support Preservation: Visualizations show FiDec refines policies along the data manifold, while isotropic methods drift into low-value or unsupported regions.
• Offline-to-Online Fine-tuning: FiDec provides a superior initialization for online fine-tuning, achieving faster convergence and higher final rewards compared to other methods.
• Ablation Studies:
  • Anisotropic vs. Isotropic: Replacing the Fisher metric with an isotropic $L_2$ penalty under the same architecture significantly degrades performance, confirming the theoretical advantage of the anisotropic approach.
  • Perturbation Time ($t_\epsilon$): The paper provides a principled analysis showing that $t_\epsilon \approx 0.8$ offers the optimal trade-off between approximation bias and numerical stability.
• Efficiency: FiDec maintains competitive training speeds (approx. 2.7ms/step), comparable to existing flow methods, as it avoids expensive iterative procedures.

5. Significance

This work fundamentally shifts the paradigm of offline RL policy refinement from global distance minimization (Wasserstein) to local geometric alignment (Fisher Information).

• Theoretical Impact: It bridges the gap between Information Geometry and Flow Matching, demonstrating that KL-constrained optimization is inherently anisotropic.
• Practical Impact: FiDec offers a simple, efficient, and highly effective solution for learning expressive policies in offline settings, particularly for tasks requiring the preservation of complex, multimodal behaviors. It resolves the "mode collapse" and "distributional shift" issues that have plagued previous flow-based approaches.