Drift Q-Learning

技術要約：Drift Q-Learning (DriftQL)

問題提起
オフライン強化学習（Offline Reinforcement Learning; RL）は、さらなる環境との相互作用なしに、固定されたデータセットから最適な方策を学習することを目的としている。この設定における中心的な課題は、「分布外（out-of-distribution; OOD）」問題である。静的なデータで訓練された価値関数は、データセットに存在しない行動に対して信頼性の低い推定値を生成してしまう。これを軽減するためには、学習された方策をデータの分布へと制約しつつ、高価値な行動へと改善することを可能にする、行動の正則化が必要となる。

拡散モデル（diffusion）やフローベースのポリシー（flow-based policies）といった表現力豊かな生成モデルを利用する既存の手法は、行動分布のマルチモーダルな構造を効果的にモデル化できる。しかし、これらは実用上の重大な制限を抱えている。具体的には、推論時に反復的なデノイジングや重いソルバーの統合を必要とするため、実行速度が遅くなる。推論を高速化するために蒸留（distillation）ベースの代替案も存在するが、追加のネットワークや2段階の訓練パイプラインを必要とするため、複雑さが増大する。

手法
著者らは、ドリフトベースの行動正則化とクリティック駆動の方策改善を組み合わせた、1ステップ生成手法である**Drift Q-Learning (DriftQL)**を提案する。DriftQLは、事前ノイズと条件付け信号からデータへの直接的な1ステップのプッシュフォワード写像を学習する「ドリフティング・モデル（Drifting Models）」のフレームワークに基づいている。

DriftQLの中核は、状態 $s$ とガウスノイズ $\epsilon$ を入力として行動 $a$ を出力する確率的アクターネットワーク $f_\theta(s, \epsilon)$ である。訓練目的関数は、以下の2つの主要なメカニズムを通じて、分布モデリングと報酬追求を統合している：

1. 条件付きドリフト場（Conditional Drift Field）: 方策は、2つの相反する力からなる状態条件付きドリフト場によって監督される：

   • 吸引（Attraction, $V^+$）: 生成された行動を観測されたデータセットの行動へと引き寄せる（行動複製：behavior cloning）。オフラインRLの設定では、通常、各状態に対して1つの観測された行動が存在するため、これは生成された行動とデータセットの行動との間の直接的なベクトル差へと簡略化される。
   • 反発（Repulsion, $V^-$）: 生成された行動同士を互いに遠ざける。これにより、方策が単一のモード（データセットの行動）に崩壊することを防ぎ、高価値な領域内での多様な候補を維持する。反発は、単一の状態に対して生成された $N$ 個の候補行動に対する、カーネル重み付き平均シフトを用いて計算される。

2. クリティック駆動の改善（Critic-Driven Improvement）: 無条件生成に使用される標準的なドリフティング・モデルとは異なり、DriftQLは価値ベースの目的関数を組み込んでいる。クリップされたダブルQクリティックが行動の価値を推定する。アクターは以下の結合損失を最小化する：
   $$L_{actor}(\theta) = \alpha L_{drift}(\theta) + L_Q(\theta)$$
   ここで、$L_{drift}$ はドリフト目標への行動回帰を表すドリフト損失であり、$L_Q$ は負の期待Q値である。価値勾配は、データのサポート（支持集合）内の高価値な領域へと方策をバイアスさせ、データ内に複数の高価値モードが存在する場合でも、それらに集中することを可能にする。

主な貢献

• 1ステップ生成方策: DriftQLは1回のフォワードパスで行動を生成するため、拡散モデルやフローベースの手法で必要とされる反復的なデノイジング・チェーン、ODEソルバー、または蒸留ステージを排除できる。
• 統一された訓練目的関数: ドリフト場による行動正則化と、Qクリティックによる方策改善を、単一のネットワークと統一された損失関数の中に統合しており、多段階の訓練パイプラインの複雑さを回避している。
• モード選択と多様性: 吸引（データへの回帰）と反発（候補間での分離）の組み合わせにより、訓練中に候補全体に質量を維持しつつ、高価値なモードへと収束させることができ、単一の決定論的な行動への崩壊や、サポートされていない領域への散逸を防ぐ。

実験結果
DriftQLは、連続制御、ナビゲーション、および操作タスクをカバーするD4RLおよびOGBenchベンチマークにおいて評価された。

• 性能: DriftQLは、78のオフライン学習タスクにおいて、最先端の拡散モデルおよびフローベースのベースライン（例：Diffusion-QL, FQL, IDQL）と同等、あるいはそれらを上回る性能を一貫して示した。特に、長期間のナビゲーションタスク（例：antmaze-large-st, antmaze-giant-st）や複雑な操作タスクにおいて顕著な向上を示した。
• 堅牢性: データ品質が低下した場合（一部の行動をランダムノイズに置き換えることでシミュレート）、DriftQLはクリーンなデータのベースラインに近い性能を維持したが、競合する手法はノイズが増加するにつれて明らかに性能が低下した。
• 効率性: 推論レイテンシに関して、DriftQLはFQLより約2倍、Diffusion-QLより3倍、IDQL/IFQLより4倍高速である。これは、別途の蒸留ステージを必要とせずに達成されている。訓練コストは、ペアワイズのカーネル計算のために生成サンプル数（$N$）に応じてスケールするが、拡散モデルにおける反復的なデノイジングのオーバーヘッドと比較すると控えめなレベルに留まっている。

意義と主張
本論文は、DriftQLを、拡散モデルやフローベースのオフラインRL手法に対する有望な代替案として位置付けている。DriftQLは、生成モデルの表現力（マルチモーダルな分布の処理）を備えつつ、決定論的なアプローチの単純さと効率性（シングルフォワードパス）を保持していると主張している。

著者らは、DriftQLがオフラインRLの相反する目的をうまく調和させていると論じている。すなわち、信頼性の高い価値推定を保証するために方策をデータのサポート内に制約しつつ、高報酬な領域への改善を許容している。反復的なサンプリングの計算オーバーヘッドを取り除くことで、DriftQLは、推論速度が極めて重要となる実世界のデプロイメントにおいて、より実用的なソリューションを提供し、同時にデータの破損に対する堅牢性も維持している。本研究は、ドリフティングがオフラインRLにとって有用な「中間領域」を提供することを示唆しており、多ステップの生成方策に伴うテスト時のコストを支払うことなく、確率的かつ分布的な訓練信号を提供できることを示している。