技術的サマリー：スコアベースのワンステップ MeanFlow 方策最適化 (SOM)

1. 問題定義

連続行動空間における強化学習 (RL) は、単一モードのガウス方策の表現力制限を克服するため、拡散モデルやフローマッチングなどの生成モデルをますます採用するようになっています。これらのモデルは、複雑で多峰性の行動分布を捉えることができ、より豊かな探索と多様な最適行動を可能にします。しかし、重要なボトルネックが残っています。標準的な拡散モデルおよびフローマッチング方策は、単一の行動を生成するために反復的な多ステップのノイズ除去を必要とします。これは推論時に莫大な計算オーバーヘッドをもたらすものであり、方策推論がデータ収集と継続的に交互に行われるオンライン RLにおいて特に有害です。

MeanFlowは、平均速度場を学習することでサンプリングを単一のネットワーク評価に削減する解決策として提案されていますが、そのオンライン RL への直接的な適用は、根本的な要件によって阻害されています。つまり、ターゲット速度場の構築には通常、ターゲット分布（最適方策）からのサンプルが必要とされるためです。オンライン RL において、これらの真のサンプルは利用できません。なぜなら、エージェントはターゲット分布を方策と同時に学習しなければならないからです。これを回避しようとする既存のアプローチ（例えば、Best-of-N サンプリング）は、行動の次元が増加するにつれてスケーラビリティの問題に直面することが多いです。

2. 手法

著者らは、明示的なターゲット分布サンプルを必要とせずに、完全なオンライン RL 環境下で MeanFlow の学習を可能にするアクター・クリティックアルゴリズムである**スコアベースのワンステップ MeanFlow 方策最適化 (SOM)**を提案します。核心的な革新は、Q 関数からのスコア推定を通じて、ターゲット速度場を直接構築することにあります。

2.1 時間依存エネルギー関数

この手法は、負の Q 関数をターゲット分布のエネルギーとして扱います。エントロピー正則化された RL フレームワークの下では、最適方策はボルツマン分布に従います。
$$\pi^*(a|s) \propto \exp(\alpha Q(s, a))$$
ここで、$\alpha$は温度パラメータです。著者らは、拡散の時間ステップ $t \in [0, 1]$ にわたってこのボルツマン構造を保持する時間依存エネルギー関数 $Q_t(s, a_t)$ を定義します。任意の時間ステップ $t$ における分布のスコアは以下のように導出されます。
$$\nabla_{a_t} \log p_t(a_t|s) = \nabla_{a_t} Q_t(s, a_t)$$
重要な点は、$Q_t$を閉形式で計算することは不可能ですが、この手法はその勾配のみを必要とし、それは近似可能であるということです。

2.2 iDEM 適応によるスコア推定

ターゲットサンプルなしでスコア $\nabla_{a_t} Q_t(s, a_t)$ を推定するために、SOM は**イテレーテッド・デノイジング・エネルギー・マッチング (iDEM)**推定量を適応させます。真のターゲット分布からサンプリングする代わりに、この手法は現在のノイズ状態 $a_t$ を中心としたガウス提案分布から候補行動をサンプリングします。これらの候補は、その Q 値に基づいて再重み付けされ、事後期待値を近似します。得られるスコア推定量は以下の通りです。
$$\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}))$$
ここで、$a^{(i)}_{0|t}$は $a_t$ に条件付けられたガウス分布から引き出されたサンプルです。これは、ノイズを通じて探索を維持しつつ、行動の更新を高価値領域へと効果的にバイアスさせます。

2.3 速度場の構築

推定されたスコアは、ターゲット速度場を導出するために**確率流 ODE (PF-ODE)**に統合されます。分散保存 (VP) SDE の場合、速度場は以下のように定義されます。
$$v_t(a_t|s) = -\frac{1}{2}\beta(t) \left( a_t + \nabla_{a_t} Q_t(s, a_t) \right)$$
この速度場は、現在の Q 関数によって誘発される高価値モードへとノイズを直接輸送します。MeanFlow 方策 $u_\theta$ は、この構築されたターゲット速度を用いて MeanFlow 恒等式を満たすように訓練され、以下の損失を最小化します。
$$\mathcal{L}(\theta) = \mathbb{E} \left[ \| u_\theta(a_t, r, t, s) - \text{sg}(u^{\text{new}}_{\text{tgt}}) \|^2_2 \right]$$
ここで、$u^{\text{new}}_{\text{tgt}}$は PF-ODE 速度から導出されます。

2.4 実装の詳細

• 正規化: スコア勾配は単位ノルムに正規化され、可変な報酬スケール全体で学習を安定させるために調整可能な係数 $w$ によって再スケーリングされます。
• クリティック学習: クリティックは、ツイン Q ネットワークを用いた分布ベルマン目的関数で更新されます。
• Best-of-N (BoN): 主な貢献はスコアベースの速度ですが、著者らはターゲット値推定のために複数のサンプルから最高 Q 値の候補を選択することで学習を加速するため、クリティック更新中に BoN サンプリングを組み込んでいます。

3. 主要な貢献

1. ボトルネックの特定: 本論文は、ターゲット分布サンプルの欠乏が、MeanFlow をオンライン RL に適用する際の主要な障壁であることを特定しました。
2. スコアベースの速度構築: 明示的なターゲットサンプルを必要とせず、Q 関数の勾配から直接ターゲット速度場を構築する新しい手法を提案しました。
3. SOM アルゴリズム: MeanFlow をオンライン RL に統合し、ワンステップの行動サンプリングで表現力豊かな方策学習を達成するアクター・クリティックアルゴリズム SOM を導入しました。
4. 実証的検証: SOM が、反復的な拡散およびフローマッチングのベースラインと比較して、移動タスクにおいて最先端のパフォーマンスを達成し、かつ学習および推論時間を大幅に削減することを示しました。

4. 実験結果

著者らは、SOM を 9 つの MuJoCo v4 連続制御タスク（Ant、HalfCheetah、Hopper、Humanoid、Walker2d、Swimmer、Pusher、Reacher、InvertedDoublePendulum を含む）で評価しました。

• パフォーマンス: SOM は、5 つの移動ベンチマークのうち 4 つで最高平均リターンを達成し、Hopper-v4 では 2 番目に高いリターンを達成しました。改善はHumanoid-v4（高次元の行動空間）で最も顕著であり、SOM は 2 番目に良い手法を明確な差で凌駕しました。これは、勾配ベースのアプローチが次元に対して有利にスケーリングすることを示唆しています。
• 効率性: 表 2 に示されるように、SOM は HalfCheetah-v4 において0.218 msのワンステップ推論時間を達成しました。これは、10〜100 ステップのノイズ除去を必要とする反復ベースラインの 0.6〜1.5 ms と比較されます。学習時間も大幅に削減されています。
• 方策分析: 可視化（t-SNE および Q 地形）は、SOM が行動サンプルを高 Q モードの周りに密に集中させるのに対し、競合するワンステップ手法（MFP など）は単一モードに崩壊するか、拡散的に分布する傾向があることを示しています。
• ロバスト性: ノイズの多い Q 関数を用いた実験は、SOM 内のガウス平滑化されたスコア推定量がロバストであり、粗いクリティックエラーを効果的に平均化しつつ、高価値モード近傍の局所構造を保持することを示しました。

5. 意義と主張

本論文は、SOM が生成方策の表現力とオンライン RL の計算制約との間の緊張関係を解決すると主張しています。ターゲット速度を Q 関数から直接導出することで、SOM はターゲットサンプルが利用不可能な完全なオンライン設定において MeanFlow 方策を訓練可能にします。

著者らは、SOM がワンステップの生成で最先端のパフォーマンスを達成し、推論遅延が重要なリアルタイム制御アプリケーションに対する実用的な解決策を提供することを強調しています。また、この手法は現在、連続制御タスク（MuJoCo）に限定され、学習されたクリティックの品質に依存していますが、効率的で表現力豊かな方策最適化に向けた重要な一歩を表しているとも述べています。今後の方向性としては、SOM を高次元または部分的に観測可能な設定（例えば、視覚制御）に拡張し、より高度な不確実性認識クリティックを統合することが挙げられます。