Résumé Technique : Optimisation de Politique MeanFlow en Une Étape Basée sur le Score (SOM)

1. Énoncé du Problème

L'apprentissage par renforcement (RL) dans des espaces d'actions continus a de plus en plus adopté des modèles génératifs, tels que la diffusion et l'appariement de flux (flow matching), pour surmonter les limitations d'expressivité des politiques gaussiennes unimodales. Ces modèles peuvent capturer des distributions d'actions complexes et multimodales, permettant une exploration plus riche et des comportements optimaux diversifiés. Cependant, un goulot d'étranglement critique subsiste : les politiques standard de diffusion et d'appariement de flux nécessitent un dénouage itératif et multi-étapes pour générer une seule action. Cela introduit une surcharge computationnelle substantielle au moment de l'inférence, ce qui est particulièrement préjudiciable en RL en ligne, où l'inférence de la politique est continuellement entrelacée avec la collecte de données.

Bien que MeanFlow ait été proposé comme solution pour réduire l'échantillonnage à une seule évaluation de réseau en apprenant un champ de vitesse moyen, son application directe au RL en ligne est entravée par une exigence fondamentale : la construction du champ de vitesse cible nécessite généralement des échantillons de la distribution cible (la politique optimale). En RL en ligne, ces échantillons de vérité terrain sont indisponibles, car l'agent doit apprendre la distribution cible simultanément à la politique. Les approches existantes tentant de contourner ce problème (par exemple, l'échantillonnage Best-of-N) ont souvent du mal à être évolutives à mesure que la dimensionnalité de l'action augmente.

2. Méthodologie

Les auteurs proposent Score-Based One-step MeanFlow Policy Optimization (SOM), un algorithme acteur-critique qui permet l'entraînement de MeanFlow dans un cadre de RL entièrement en ligne sans nécessiter d'échantillons explicites de la distribution cible. L'innovation centrale réside dans la construction du champ de vitesse cible directement à partir de la fonction Q via l'estimation du score.

2.1 Fonction d'Énergie Dépendante du Temps

La méthode traite la fonction Q négative comme l'énergie de la distribution cible. Dans un cadre de RL régularisé par l'entropie, la politique optimale suit une distribution de Boltzmann :
$$\pi^*(a|s) \propto \exp(\alpha Q(s, a))$$
où $\alpha$ est un paramètre de température. Les auteurs définissent une fonction d'énergie dépendante du temps $Q_t(s, a_t)$ qui préserve cette structure de Boltzmann à travers les pas de temps de diffusion $t \in [0, 1]$. Le score de la distribution à n'importe quel pas de temps $t$ est dérivé comme suit :
$$\nabla_{a_t} \log p_t(a_t|s) = \nabla_{a_t} Q_t(s, a_t)$$
Crucialement, bien que le calcul de $Q_t$ sous forme close soit intraitable, la méthode ne nécessite que son gradient, qui peut être approché.

2.2 Estimation du Score via Adaptation de iDEM

Pour estimer le score $\nabla_{a_t} Q_t(s, a_t)$ sans échantillons cibles, SOM adapte l'estimateur Iterated Denoising Energy Matching (iDEM). Au lieu d'échantillonner à partir de la véritable distribution cible, la méthode échantillonne des actions candidates à partir d'une proposition gaussienne centrée sur l'état bruité actuel $a_t$. Ces candidats sont repondérés en fonction de leurs valeurs Q pour approximer l'espérance a posteriori. L'estimateur de score résultant est :
$$\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}))$$
où $a^{(i)}_{0|t}$ sont des échantillons tirés d'une distribution gaussienne conditionnée par $a_t$. Cela biaise efficacement les mises à jour d'actions vers des régions à haute valeur tout en maintenant l'exploration grâce au bruit.

2.3 Construction du Champ de Vitesse

Le score estimé est intégré dans une Équation Différentielle Ordinaire de Flux de Probabilité (PF-ODE) pour dériver le champ de vitesse cible. Pour une Équation Différentielle Stochastique (SDE) à Préservation de Variance (VP), le champ de vitesse est défini comme :
$$v_t(a_t|s) = -\frac{1}{2}\beta(t) \left( a_t + \nabla_{a_t} Q_t(s, a_t) \right)$$
Ce champ de vitesse transporte le bruit directement vers les modes à haute valeur induits par la fonction Q actuelle. La politique MeanFlow $u_\theta$ est ensuite entraînée pour satisfaire l'identité MeanFlow en utilisant ce champ cible construit, en minimisant la perte :
$$\mathcal{L}(\theta) = \mathbb{E} \left[ \| u_\theta(a_t, r, t, s) - \text{sg}(u^{\text{new}}_{\text{tgt}}) \|^2_2 \right]$$
où $u^{\text{new}}_{\text{tgt}}$ est dérivé de la vitesse PF-ODE.

2.4 Détails de l'Implémentation

• Normalisation : Le gradient du score est normalisé à une norme unitaire et redimensionné par un coefficient ajustable $w$ pour stabiliser l'entraînement à travers différentes échelles de récompense.
• Entraînement du Critique : Le critique est mis à jour en utilisant un objectif de Bellman distributionnel avec des réseaux Q jumeaux.
• Best-of-N (BoN) : Bien que la contribution principale soit la vitesse basée sur le score, les auteurs intègrent l'échantillonnage BoN lors des mises à jour du critique pour accélérer l'entraînement en sélectionnant le candidat à Q le plus élevé parmi plusieurs échantillons pour l'estimation de la valeur cible.

3. Contributions Clés

1. Identification d'un Goulot d'Étranglement : L'article identifie le manque d'échantillons de la distribution cible comme la principale barrière à l'application de MeanFlow en RL en ligne.
2. Construction de Vitesse Basée sur le Score : Il propose une méthode novatrice pour construire le champ de vitesse cible directement à partir des gradients de la fonction Q, contournant ainsi le besoin d'échantillons cibles explicites.
3. Algorithme SOM : L'introduction de SOM, un algorithme acteur-critique qui intègre MeanFlow avec le RL en ligne, réalisant un apprentissage de politique expressif avec un échantillonnage d'action en une seule étape.
4. Validation Empirique : Démonstration que SOM atteint des performances de pointe sur des tâches de locomotion tout en réduisant considérablement les temps d'entraînement et d'inférence par rapport aux bases de diffusion itérative et d'appariement de flux.

4. Résultats Expérimentaux

Les auteurs ont évalué SOM sur neuf tâches de contrôle continu MuJoCo v4 (incluant Ant, HalfCheetah, Hopper, Humanoid, Walker2d, Swimmer, Pusher, Reacher et InvertedDoublePendulum).

• Performance : SOM a obtenu le retour moyen le plus élevé sur quatre des cinq benchmarks de locomotion et le deuxième plus élevé sur Hopper-v4. L'amélioration a été la plus prononcée sur Humanoid-v4 (haute dimensionnalité d'action), où SOM a surpassé la deuxième meilleure méthode avec une marge claire, suggérant que l'approche basée sur le gradient s'adapte favorablement à la dimension.
• Efficacité : Comme indiqué dans le Tableau 2, SOM atteint un temps d'inférence en une seule étape de 0,218 ms sur HalfCheetah-v4, contre 0,6–1,5 ms pour les bases itératives (qui nécessitent 10–100 étapes de débruitage). Le temps d'entraînement est également considérablement réduit.
• Analyse de la Politique : Les visualisations (t-SNE et paysages Q) indiquent que SOM concentre les échantillons d'actions étroitement autour des modes à haute Q, tandis que les méthodes concurrentes en une seule étape (comme MFP) tendent à s'effondrer vers des modes uniques ou à se distribuer de manière diffuse.
• Robustesse : Les expériences avec des fonctions Q bruitées ont démontré que l'estimateur de score lissé par une gaussienne dans SOM reste robuste, moyennant efficacement les erreurs grossières du critique tout en préservant la structure locale près des modes à haute valeur.

5. Signification et Revendications

L'article revendique que SOM résout la tension entre l'expressivité des politiques génératives et les contraintes computationnelles du RL en ligne. En dérivant la vitesse cible directement à partir de la fonction Q, SOM permet aux politiques MeanFlow d'être entraînées dans des contextes entièrement en ligne où les échantillons cibles sont indisponibles.

Les auteurs soulignent que SOM atteint des performances de pointe avec une étape de génération unique, offrant une solution pratique pour les applications de contrôle en temps réel où la latence d'inférence est critique. Ils notent que bien que la méthode soit actuellement limitée aux tâches de contrôle continu (MuJoCo) et dépende de la qualité du critique appris, elle représente une étape significative vers une optimisation de politique efficace et expressive. Le travail suggère que les directions futures incluent l'extension de SOM à des contextes de plus haute dimension ou partiellement observables (par exemple, le contrôle visuel) et l'intégration de critiques plus avancés conscients de l'incertitude.