Technical Summary: Score-Based One-step MeanFlow Policy Optimization (SOM)

1. Problem Statement

Reinforcement learning (RL) in continuous action spaces has increasingly adopted generative models, such as diffusion and flow matching, to overcome the expressivity limitations of unimodal Gaussian policies. These models can capture complex, multimodal action distributions, enabling richer exploration and diverse optimal behaviors. However, a critical bottleneck remains: standard diffusion and flow-matching policies require iterative, multi-step denoising to generate a single action. This introduces substantial computational overhead at inference time, which is particularly detrimental in online RL, where policy inference is continuously interleaved with data collection.

While MeanFlow has been proposed as a solution to reduce sampling to a single network evaluation by learning an average velocity field, its direct application to online RL is hindered by a fundamental requirement: constructing the target velocity field typically necessitates samples from the target distribution (the optimal policy). In online RL, these ground-truth samples are unavailable, as the agent must learn the target distribution simultaneously with the policy. Existing approaches that attempt to bypass this (e.g., Best-of-N sampling) often struggle with scalability as action dimensionality increases.

2. Methodology

The authors propose Score-Based One-step MeanFlow Policy Optimization (SOM), an actor-critic algorithm that enables MeanFlow training in fully online RL without requiring explicit target-distribution samples. The core innovation lies in constructing the target velocity field directly from the Q-function via score estimation.

2.1 Time-Dependent Energy Function

The method treats the negative Q-function as the energy of the target distribution. Under an entropy-regularized RL framework, the optimal policy follows a Boltzmann distribution:
$$\pi^*(a|s) \propto \exp(\alpha Q(s, a))$$
where $\alpha$ is a temperature parameter. The authors define a time-dependent energy function $Q_t(s, a_t)$ that preserves this Boltzmann structure across the diffusion timesteps $t \in [0, 1]$. The score of the distribution at any timestep $t$ is derived as:
$$\nabla_{a_t} \log p_t(a_t|s) = \nabla_{a_t} Q_t(s, a_t)$$
Crucially, while computing $Q_t$ in closed form is intractable, the method only requires its gradient, which can be approximated.

2.2 Score Estimation via iDEM Adaptation

To estimate the score $\nabla_{a_t} Q_t(s, a_t)$ without target samples, SOM adapts the Iterated Denoising Energy Matching (iDEM) estimator. Instead of sampling from the true target distribution, the method samples candidate actions from a Gaussian proposal centered at the current noisy state $a_t$. These candidates are reweighted based on their Q-values to approximate the posterior expectation. The resulting score estimator is:
$$\nabla_{a_t} Q_t(s, a_t) \approx \nabla_{a_t} \log \sum_{i=1}^K \exp(\alpha Q(s, a^{(i)}_{0|t}))$$
where $a^{(i)}_{0|t}$ are samples drawn from a Gaussian distribution conditioned on $a_t$. This effectively biases action updates toward high-value regions while maintaining exploration through the noise.

2.3 Velocity Field Construction

The estimated score is integrated into a Probability Flow ODE (PF-ODE) to derive the target velocity field. For a Variance Preserving (VP) SDE, the velocity field is defined as:
$$v_t(a_t|s) = -\frac{1}{2}\beta(t) \left( a_t + \nabla_{a_t} Q_t(s, a_t) \right)$$
This velocity field transports noise directly toward high-value modes induced by the current Q-function. The MeanFlow policy $u_\theta$ is then trained to satisfy the MeanFlow identity using this constructed target velocity, minimizing the loss:
$$\mathcal{L}(\theta) = \mathbb{E} \left[ \| u_\theta(a_t, r, t, s) - \text{sg}(u^{\text{new}}_{\text{tgt}}) \|^2_2 \right]$$
where $u^{\text{new}}_{\text{tgt}}$ is derived from the PF-ODE velocity.

2.4 Implementation Details

• Normalization: The score gradient is normalized to unit norm and rescaled by a tunable coefficient $w$ to stabilize training across varying reward scales.
• Critic Training: The critic is updated using a distributional Bellman objective with twin Q-networks.
• Best-of-N (BoN): While the primary contribution is the score-based velocity, the authors incorporate BoN sampling during critic updates to accelerate training by selecting the highest-Q candidate from multiple samples for the target value estimation.

3. Key Contributions

1. Identification of a Bottleneck: The paper identifies the lack of target-distribution samples as the primary barrier to applying MeanFlow in online RL.
2. Score-Based Velocity Construction: It proposes a novel method to construct the target velocity field directly from Q-function gradients, circumventing the need for explicit target samples.
3. SOM Algorithm: The introduction of SOM, an actor-critic algorithm that integrates MeanFlow with online RL, achieving expressive policy learning with single-step action sampling.
4. Empirical Validation: Demonstration that SOM achieves state-of-the-art performance on locomotion tasks while significantly reducing training and inference time compared to iterative diffusion and flow-matching baselines.

4. Experimental Results

The authors evaluated SOM on nine MuJoCo v4 continuous control tasks (including Ant, HalfCheetah, Hopper, Humanoid, Walker2d, Swimmer, Pusher, Reacher, and InvertedDoublePendulum).

• Performance: SOM achieved the highest mean return on four out of five locomotion benchmarks and the second-highest on Hopper-v4. The improvement was most pronounced on Humanoid-v4 (high action dimensionality), where SOM outperformed the second-best method by a clear margin, suggesting the gradient-based approach scales favorably with dimension.
• Efficiency: As shown in Table 2, SOM achieves a single-step inference time of 0.218 ms on HalfCheetah-v4, compared to 0.6–1.5 ms for iterative baselines (which require 10–100 denoising steps). Training time is also substantially reduced.
• Policy Analysis: Visualizations (t-SNE and Q-landscapes) indicate that SOM concentrates action samples tightly around high-Q modes, whereas competing one-step methods (like MFP) tend to collapse to single modes or distribute diffusely.
• Robustness: Experiments with noisy Q-functions demonstrated that the Gaussian-smoothed score estimator in SOM remains robust, effectively averaging out coarse critic errors while preserving local structure near high-value modes.

5. Significance and Claims

The paper claims that SOM resolves the tension between the expressivity of generative policies and the computational constraints of online RL. By deriving the target velocity directly from the Q-function, SOM enables MeanFlow policies to be trained in fully online settings where target samples are unavailable.

The authors emphasize that SOM achieves state-of-the-art performance with a single generation step, offering a practical solution for real-time control applications where inference latency is critical. They note that while the method is currently limited to continuous control tasks (MuJoCo) and relies on the quality of the learned critic, it represents a significant step toward efficient, expressive policy optimization. The work suggests that future directions include extending SOM to higher-dimensional or partially observable settings (e.g., visual control) and integrating more advanced uncertainty-aware critics.