Technical Summary: Entropy-Regularized Adjoint Matching for Offline RL

Problem Statement

The paper addresses the Support-Binding Dilemma (SBD) inherent in offline Reinforcement Learning (RL) when utilizing expressive generative policies, such as flow-matching models. While continuous-time frameworks like Adjoint Matching (QAM) stabilize policy optimization by bypassing unstable solver backpropagation, they remain strictly bound to the fixed behavior distribution ($\pi_\beta$) of the offline dataset. This dependence creates two critical limitations:

1. Popularity Bias (Density Front): The optimization relies on a multiplicative dependence on the behavior prior ($\pi^* \propto \pi_\beta \exp(Q/\tau)$). If an optimal action is rare in the dataset, the vanishing $\pi_\beta$ term suppresses the reward signal. Overcoming this requires aggressive temperature scaling, which amplifies critic approximation errors and forces the agent to accept sub-optimal, high-frequency actions to maintain stability.
2. Zero-Support Trap (Geometric Front): If an optimal action lies outside the support of the behavior data, $\pi_\beta = 0$. This structural constraint prevents the agent from ever reaching these regions, halting the discovery of novel behaviors regardless of the reward signal.

Existing workarounds, such as appending residual Gaussian policies (e.g., QAM-Edit) to force the agent off the data manifold, often break the continuous Ordinary Differential Equation (ODE) formulation. These heuristics act as structural patches that disrupt the coherent vector field trajectory and reintroduce the expressivity bottlenecks associated with unimodal distributions.

Methodology: Maximum Entropy Adjoint Matching (ME-AM)

The authors propose Maximum Entropy Adjoint Matching (ME-AM), a unified framework that resolves the SBD entirely within the continuous flow formulation. ME-AM integrates two primary mechanisms:

1. Reward-Guided Mixture Prior (Geometric Expansion)

To bridge the geometric gap between isolated data clusters without breaking the continuous ODE, ME-AM augments the behavior dataset during training.

• Mechanism: A fraction $\epsilon$ of the dataset actions is replaced with synthetic targets generated by an auxiliary Gaussian actor ($\pi_\omega$). This actor is trained via a Soft Actor-Critic (SAC) style objective to predict high-reward actions conditioned on the state and an empirical action.
• Effect: The base flow is trained against a mixture distribution $\pi_{mix} = (1-\epsilon)\pi_\beta + \epsilon\pi_\omega$. This natively extends the valid geometric support, allowing the model to learn to transport noise to out-of-distribution (OOD) modes while maintaining a structurally sound, continuous ODE.

2. Density Flattening via Mirror Descent (Entropy Maximization)

To mitigate popularity bias, ME-AM maximizes the differential entropy of the fine-tuned policy directly within the adjoint system.

• Mechanism: Instead of explicit density evaluation (which is computationally prohibitive), the method employs Functional Mirror Descent (MD). The functional gradient for entropy maximization simplifies to the negative score function ($-\nabla \log \pi$).
• Implementation: The negative score is injected directly into the terminal boundary condition of the Adjoint ODE alongside the value gradient. To avoid numerical instability at the terminal boundary ($t \to 1$), a lightweight auxiliary score network ($S_\psi$) is trained via a variance-weighted denoising objective.
• Effect: This process flattens the peaked behavior density across the extended support, allowing the agent to emphasize high-reward actions without requiring the extreme, error-amplifying contrast of strict temperature scaling.

Theoretical Formulation

The unified objective aligns the fine-tuned vector field with an interpolated reference flow ($v_{ref}$) that balances the static mixture base and a slow-moving target anchor. The authors prove (Theorem 2) that the optimal stationary policy converges to:
$$\pi^*(a|s) \propto (\pi_{mix}(a|s))^{\delta(s)} \exp\left( \frac{1}{\kappa(s)} Q_\phi(s, a) \right)$$
where $\delta(s) \in (0, 1)$ is a state-dependent exponent that mathematically guarantees the flattening of the prior density, decoupling exploration from empirical data frequencies.

Key Contributions

1. Unified Framework: ME-AM is the first method to solve the Support-Binding Dilemma entirely within a continuous flow formulation, eliminating the need for disjointed edits or post-processing heuristics that break the ODE structure.
2. Geometric and Density Resolution: It simultaneously addresses the zero-support trap via a learnable, critic-guided mixture prior and the popularity bias via functional Mirror Descent entropy maximization.
3. Theoretical Guarantees: The paper provides proofs establishing that the continuous-time objective converges exactly to a targeted, density-flattened stationary policy.
4. Robust Implementation: The method utilizes a decoupled auxiliary score parameterization to bypass the numerical explosion of analytical scores at the terminal boundary, ensuring stable training.

Experimental Results

The authors evaluated ME-AM on two challenging domains from OGBench (puzzle-4x4 and cube-triple) characterized by sparse rewards and complex, multi-modal behaviors.

• Offline Performance: ME-AM demonstrated competitive or superior performance compared to a diverse suite of State-of-the-Art (SOTA) methods, including QAM, FEdit, CGQL, and ReBRAC. On the puzzle-4x4 domain, ME-AM achieved a global success rate of 56% (aggregate), significantly outperforming the next best method (QAM-E at 37%).
• Ablation Studies:
  • Geometric Expansion: Relying solely on behavior support ($\epsilon=0$) caused total performance collapse in sparse domains, confirming the necessity of the mixture prior.
  • KL Regularization: Explicit KL path constraints were found essential for stabilizing training when stochasticity was introduced during geometric expansion.
  • Noise Floor: A larger evaluation noise floor ($\sigma_{min} = 0.3$) improved performance in certain tasks, acting as a spatial smoothing factor.
• Offline-to-Online Fine-Tuning: ME-AM significantly improved online performance compared to vanilla QAM, achieving results on par with QAM-edit while eliminating the need for an additional edit policy for online exploration.
• Efficiency: While training time is slightly higher due to auxiliary networks, the inference speed (0.14 ms/step) matches standard flow-matching policies, as the auxiliary components are severed from the computational graph prior to deployment.

Significance and Claims

The paper claims that ME-AM represents a structural advancement in offline RL by resolving the fundamental tension between maximizing reward and avoiding out-of-distribution extrapolation. By extending the valid action manifold and flattening the empirical behavioral density within a single, coherent continuous-time framework, ME-AM enables agents to extract robust, optimal policies from sparse, offline datasets without the instability or expressivity limitations of previous heuristic approaches. The authors emphasize that this approach allows for the discovery of novel behaviors in regions devoid of data support while maintaining the mathematical guarantees of continuous optimal control.