テクニカルサマリー：オフライン強化学習における拡散方策のためのノイジー空間方策勾配（Noisy-Space Policy Gradient）

1. 問題提起

オフライン強化学習（Offline Reinforcement Learning; RL）は、環境とのさらなる相互作用なしに、固定されたデータセットから効果的な方策を学習することを目的としている。この設定における主要な課題は「分布シフト（distributional shift）」問題である。データセットのサポートを超えて外挿する方策は、しばしばバイアスのかかった価値推定を引き起こす。このような複雑でマルチモーダルな行動分布に対処するため、拡散方達（diffusion policies）のような表現力豊かな生成モデルが登場している。拡散方策は、ノイズの乗った潜在変数から最終的なクリーンな行動へと、反復的なデノイジング過程を通じて行動を生成する。

しかし、拡散方策を価値ベースのRLと統合することには、根本的な構造的不一致が存在する：

• 最適化空間： 方策の最適化は、ノイジーな潜在空間（デノイジング過程における変数 $x^{(k)}$ のシーケンス）で行われる。
• 価値の定義： 強化学習の目的関数（価値関数、アドバンテージ、勾配）は、実行されるクリーンな行動（$a \in \mathcal{A}$）に対して定義される。

このギャップを埋める既存のアプローチには、以下のような重大な欠点がある：

1. 行動空間でのサンプリング： 拡散モデルを単なるサンプラーとして扱い、行動空間のクリティックによって行動を再重み付けする方法は、潜在空間でのガイダンスをバイパスしてしまう。
2. エンドツーエンドのバックプロパゲーション： デノイジング・トラジェトリ全体を通じて価値勾配をバックプロパゲートすることで方策を最適化すると、長い確率的計算グラフが作成され、計算コストの増大と不安定化を招く。
3. 直接的な潜在変数評価： 中間変数に対してQ関数を評価することで、ノイジーな潜在変数に直接価値ガイダンスを注入する方法は、それらの変数が報酬を決定する行動ではないという事実を無視している。

本論文は、ノイジーな潜在空間の最適化と行動空間の報酬を、原理的かつ意味的に一貫した方法で橋渡しする価値定式化の欠如を指摘している。

2. 手法

著者らは、**ノイジー空間行動価値関数（Noisy-Space Action-Value Function）と、そこから導出されるノイジー空間方策勾配（Noisy-Space Policy Gradient; NSPG）**を中心としたフレームワークを提案している。

2. 1 ノイジー空間行動価値関数 ($Q_{noisy}$)

核心となる理論的貢献は、ノイズ自体に固有の価値を割り当てるのではなく、それが誘発するクリーンな行動の分布に対して価値を割り当てる、ノイジーな潜在変数のための価値関数の定義である。

• 定義： 状態 $s$、拡散ステップ $k$、およびノイジーな潜在変数 $x^{(k)}$ に対して、ノイジー空間行動価値関数は、$(s, x^{(k)}, k)$ に条件付けられた拡散方策による期待行動空間リターンとして次のように定義される：
  $$Q_{noisy}^\pi(s, x^{(k)}, k) \triangleq \mathbb{E}_{a \sim \pi(\cdot|s, x^{(k)}, k)} [Q^\pi(s, a)]$$
• 意味論的正当性： 著者らは、$Q_{noisy}$ が、方策が $(s, x^{(k)}, k)$ で初期化され、環境内で実行されたときに得られる期待割引リターンと厳密に一致することを証明している。これにより、環境インターフェース（クリーンな行動のみを観測する）と価値関数が整合していることが保証される。
• サロゲート意味論： 方策の逆時間ダイナミクスを変更せずに条件付けを適切に定義するため、著者らはサロゲート結合分布を導入している。これは、クリーンな行動を腐敗させる順方向の拡散プロセスを通じて、クリーンな行動 $a$ とノイジーな潜在変数 $x^{(k)}$ の関係をモデル化するものである。これにより、生成メカニズムを変更することなく勾配を導出することが可能になる。

2. 2 ノイジー空間方策勾配 (NSPG)

サロゲート意味論を用いて、著者らは $x^{(k)}$ に関する $Q_{noisy}$ の勾配の閉形式表現を導出した。

• 勾配の定式化：
  $$\nabla_{x^{(k)}} Q_{noisy}^\pi(s, x^{(k)}, k) = \mathbb{E}_{a \sim \pi} \left[ A_{noisy}^\pi(s, a, x^{(k)}, k) \nabla_{x^{(k)}} \log q_k(x^{(k)} | a) \right]$$
  ここで、
  • $A_{noisy}^\pi = Q^\pi(s, a) - Q_{noisy}^\pi(s, x^{(k)}, k)$ は、ノイジー空間のアドバンテージである。
  • $\nabla_{x^{(k)}} \log q_k(x^{(k)} | a)$ は、順方向拡散カーネルのスコア関数である。
• 重要な相違点： ノイジーな潜在変数を単一のデコードされた行動 $\hat{a}$ にマップして $Q(s, \hat{a})$ を微分する手法とは異なり、NSPGは適合するクリーンな行動の分布に対する期待値を微分する。これにより、勾配は、最も確率の高い行動が最高値を持っていなくても、高価値な適合行動に関連するノイジー空間内の方向を強化し、低価値な方向を抑制することができる。
• バックプロパゲーションの回避： 決定的なのは、この勾配がクリーンな行動空間の価値推定（$Q^\pi$）と順方向拡散カーネルのみを使用して計算される点である。これにより、フルデノイジング・トラジェトリへのバックプロパゲーションを回避でき、安定性と効率性が大幅に向上する。

2. 3 ノイジー空間アクター・クリティック (NSAC)

著者らは、このフレームワークを実用的なオフラインRLアルゴリズムである NSAC として実装している：

1. クリティック： 過大評価と分布シフトを軽減するために、オフラインデータと下限信頼境界（LCB）ターゲットを用いた保守的な行動空間Q関数（$Q_\phi$）を学習する。クリティックは、ノイジーな潜在変数に対して決して評価されない。
2. アクター： 標準的なデノイジング損失（行動模倣）と、NSPGから導出された価値ガイダンス項を組み合わせた回帰目的関数を用いて、拡散方策を更新する。
   • アクターの損失には、$\lambda \cdot \epsilon_\theta^\top \nabla_{x^{(k)}} Q_{noisy}$ という項が含まれる。
   • これにより、方策はデータ分布に留まりながら、潜在空間において改善することが可能となる。

3. 主な貢献

1. 不一致の特定： 本論文は、拡散ベースのRLにおける、ノイジーな潜在空間の最適化と行動空間の価値関数の間の意味的なギャップを形式的に特定した。
2. 理論的定式化： $Q_{noisy}$ の導入。これは、クリーンな行動が生成する分布を通じてのみノイジーな潜在変数に価値を割り当てるものであり、MDPとの意味的一貫性を保証する。
3. NSPGの導出： クリーンな行動空間の価値推定のみを使用する方策勾配を導出し、デノイジングプロセスへのバックプロパゲーションを不要にした。
4. NSACアルゴリズム： NSPGを実装した安定したアクター・クリティック・アルゴリズムの開発。これにより、表現力豊かな拡散方策をオフライン設定で効果的に訓練できることを示した。

4. 実験結果

著者らは、状態ベースの D4RL ベンチマーク（Locomotion, AntMaze, Adroit）およびビジョンベースの OGBench タスクを用いてNSACを評価している。

• パフォーマンス： NSACは、多様な環境において最先端または競争力のある性能を達成している。
  • Locomotion： HalfCheetah, Hopper, Walker2dにおいて、先行手法（例：CQL, IQL, Diffuser）と同等またはそれを上回る性能を示す。特に、HalfCheetah medium-replayにおいて約10%の向上を達成している。
  • AntMaze： 既存手法に対し、特に大規模な迷路環境において大幅な改善を示し、最も近いベースラインであるDACに対し相対的に50%以上の性能向上を実現した。
  • Adroit： 評価された手法の中で、Pen-Cloned-v1において最強の性能を達成した。
  • Vision (OGBench)： ビジョンタスクにおいて、フローベースの方策やDACを上回る最高の総スコア（271.5）を獲得した。
• アブレーション研究：
  • ガイダンス強度 ($\lambda$)： 本手法は、中間的な値が最適なパフォーマンスをもたらす範囲の、幅広いガイダンス強度に対して頑健である。過度なガイダンスは、方策をデータのサポートから押し出すため、性能を低下させる。
  • モンテカルロ近似： ノイジー空間の期待値は、少数のサンプル（例：2〜4個）を用いて効率的に推定でき、DACと比較して計算オーバーヘッドは最小限である。
  • 悲観性（Pessimism）： アルゴリズムは、広範な悲観性係数の範囲において頑健である。

5. 意義と主張

本論文は、NSPGが拡散方策を訓練するための原理的かつ効果的な基盤を提供すると主張している。ノイジーな潜在空間の最適化と行動空間の価値との間の不一致を解決することで、本手法はエンドツーエンドのバックプロパゲーションの不安定さや、ノイジーな変数に対してクリティックを評価することの意味的な不整合を回避している。

著者らは、彼らのアプローチが以下の特性を持つことを強調している：

• 扱いにくさを犠牲にすることなく、マルチステップ拡散方策の表現力を維持する。
• 価値推定と方策最適化が密接に結合している、困難で長期的なホライゾンを持つマルチモーダルな設定において、控えめながらも重要な安定性と性能の向上を提供する。
• このフレームワークはオンラインおよびオフライン・トゥ・オンライン学習にも自然に拡張できる可能性があるが、これは今後の課題として残されている。

本研究は、すべてのオフラインRL問題を解決すると主張するものではなく、特に拡散方策と価値ベースの目的関数の統合という問題に焦点を当て、理論的に裏付けられた、ヒューリスティックなサンプリングや不安定なエンドツーエンドの訓練に代わる選択肢を提示している。