Technical Summary: Stochastic MeanFlow Policies (SMFP)

1. Problem Statement

Online off-policy reinforcement learning (RL) for high-dimensional continuous control faces a fundamental trade-off between policy expressivity and optimisation tractability.

• Policy Representation: Standard Gaussian policies offer fast inference and tractable entropy estimation but fail to model complex, multimodal action distributions. Conversely, generative policies (e.g., diffusion or flow-based) provide rich multimodal expressivity but suffer from intractable likelihoods, difficult entropy estimation, and high inference latency due to iterative sampling.
• Optimisation Methodology: There is a tension between entropy-regularised exploration (e.g., Soft Actor-Critic, SAC), which promotes broad exploration but can lead to instability, and mirror descent (MD) updates, which ensure stable, conservative policy improvements by constraining updates to remain close to the previous policy.
• The Core Conflict: Combining SAC-style entropy maximisation with MD constraints creates a multimodal target distribution (a composite of a value-induced Boltzmann distribution and the previous policy). Standard unimodal Gaussian policies cannot represent this target, leading to a representation mismatch. While expressive generative policies could theoretically bridge this gap, they typically lack a tractable mechanism for explicit entropy control, making the combination of entropy-driven exploration and mirror-descent stability practically infeasible.

2. Methodology: Stochastic MeanFlow Policies (SMFP)

The authors propose Stochastic MeanFlow Policies (SMFP), a one-step generative policy framework designed to resolve the above conflicts by combining MeanFlow-based noise-to-action mappings with Gaussian reparameterisation.

2.1 Stochastic Reparameterisation

SMFP reformulates the deterministic MeanFlow identity into a stochastic mapping. Instead of a purely deterministic velocity field, the policy $g_\theta$ predicts both a mean action and a conditional standard deviation:
$$g_\theta(a_t, b, t) = \underbrace{a_t - u_\theta(a_t, b, t)}_{\text{Mean } \mu} + \underbrace{\sigma_\theta(a_t, b, t) \odot \epsilon}_{\text{Stochasticity } \sigma}$$
where $a_t$ is the interpolation between noise and action, $u_\theta$ is the learned average velocity, and $\epsilon$ is standard Gaussian noise. This structure allows the policy to output actions in a single forward pass (one-step inference) while maintaining a conditional Gaussian distribution.

2.2 Tractable Entropy Surrogate

A key innovation is the derivation of a tractable entropy surrogate. Since the marginal entropy of the generative policy is intractable, SMFP maximises a lower bound based on the conditional entropy given the latent noise $e$.

• The objective uses the expected log-standard deviation of the predicted Gaussian: $\mathbb{E}[\sum \log \sigma^{(i)}_\theta]$.
• To prevent variance collapse without driving unbounded entropy growth, a hinge-style entropy regulariser is applied. It activates only when the normalised log-standard deviation falls below a threshold $\kappa_\sigma$, ensuring a minimum noise scale for exploration.

2.3 Off-Policy Mirror Descent Optimisation

SMFP integrates this entropy control with off-policy mirror descent. The optimisation objective combines:

1. Value Maximisation: Standard Q-learning loss.
2. Entropy Regularisation: The tractable conditional entropy surrogate described above.
3. Mirror Descent Constraint: A KL-divergence term anchoring the update to the previous policy $\pi_{\theta_{old}}$.

Since the target distribution for mirror descent is non-parametric and unnormalised, SMFP employs a weighted regression approach. It samples actions from the previous policy and weights them by an advantage-based term (approximated as $\max(0, Q(s,a) - V(s))$) to project the generative policy onto the optimal target. This avoids the need for explicit density evaluation of the target.

2.4 Value-Guided Action Selection

To further reduce variance, SMFP utilizes Best-of-K sampling. For a given state, $K$ action candidates are generated in parallel (using JAX vmap), and the action with the highest Q-value is selected. This acts as a lightweight rejection filter with negligible inference overhead.

3. Key Contributions

1. SMFP Architecture: The introduction of a one-step stochastic generative policy that combines MeanFlow expressivity with Gaussian reparameterisation, enabling single-step inference while supporting explicit entropy control.
2. Tractable Entropy Control: The derivation of a conditional entropy lower bound that makes entropy-regularised mirror-descent RL practically effective for the first time with generative policies.
3. Unified Objective: A unified loss function that jointly optimises entropy-driven exploration, value-weighted mirror descent, and one-step generative inference, resolving the representation mismatch between multimodal targets and unimodal policies.
4. Empirical Validation: Demonstration that combining SAC-style exploration with mirror descent is beneficial when supported by expressive generative policies, contrary to findings in Gaussian-only settings.

4. Experimental Results

The authors evaluated SMFP on seven MuJoCo continuous control benchmarks (Hopper, Walker2D, Ant, HalfCheetah, Humanoid, HumanoidStandup, Swimmer) over 30 random seeds.

• Performance: SMFP achieved the best aggregate performance across all benchmarks, outperforming both classical Gaussian baselines (PPO, TD3, SAC) and state-of-the-art generative baselines (Diffusion-based like DIME, QVPO; Flow-based like SAC Flow, FPMD).
• Efficiency: Unlike iterative diffusion policies that require 16–20 sampling steps, SMFP maintains single-step inference (NFE=1). On an NVIDIA RTX 5090, SMFP achieved strong performance with negligible inference latency (0.42ms), comparable to standard Gaussian policies and significantly faster than iterative generative methods.
• Ablation Studies:
  • Removing the entropy regulariser degraded performance, confirming the need for stochastic exploration.
  • Removing the mirror descent term reduced stability, particularly in high-dimensional tasks like Humanoid.
  • The method showed robustness to hyperparameter variations in the entropy temperature ($\alpha$) and mirror descent coefficient ($\lambda$).

5. Significance and Claims

The paper claims that SMFP represents a significant step forward in online off-policy RL by successfully bridging the gap between expressive generative modelling and stable, entropy-regularised optimisation.

• Practical Effectiveness: The authors state that SMFP makes "entropy-regularised mirror-descent RL practically effective for the first time." Previous attempts to combine these objectives failed under Gaussian parameterisations due to the inability to represent multimodal targets; SMFP overcomes this via its generative nature.
• Balanced Trade-off: The work demonstrates that it is possible to retain the multimodal expressivity of generative policies and the stability of mirror descent without sacrificing the inference speed required for real-time control.
• Limitations: The authors modestly note that the tractable surrogate relies specifically on the reparameterised structure of MeanFlow. Extending this to other generative classes may require new entropy bounds. Additionally, the current implementation uses fixed hyperparameters for $\alpha$ and $\lambda$, leaving adaptive weighting schemes as future work.

In summary, SMFP provides a principled framework where the complementary advantages of expressive generative policies and efficient Gaussian-based entropy control are preserved, enabling stable and high-performing online RL in complex continuous domains.