技術要約: Lagrangian Perturbation Diffusion Steering (LP-DS)

1. 問題提起

拡散モデルやフロー・マッチングに基づく高容量の生成方策は、連続制御およびロボティクスにおける強力なパラダイムとして確立されています。これらのモデルは、オフラインデータから学習された、表現力豊かでマルチモーダルな行動分布を表現することに長けています。しかし、これらは主に2つの制限に直面しています：

1. 行動クローニング（BC）の制約: BCを通じて訓練された方策は、不完全な状態空間の被覆、稀なリカバリー動作の欠落、および分布シフトに対する脆弱性を含む、デモンストレーション・データセットの限界をそのまま引き継いでしまいます。
2. 強化学習（RL）の不安定性: RLはオンラインでの改善メカニズムを提供しますが、大規模な生成デコーダを直接ファインチューニングすることは、しばしば不安定でサンプル効率が悪くなります。最適化目的関数と、長期間のデノイジングまたは積分ダイナミクスとの相互作用が、学習の不安定性を招くことがよくあります。

Diffusion Steering via Reinforcement Learning (DSRL) のような既存の代替手法は、凍結された方策を改良するために潜在ノイズ空間の方策を学習しようと試みています。しかし、DSRLは事前学習された事前分布を学習された潜在方策によって事実上置き換えてしまいます。このアプローチには、以下の2つの失敗モードがあります：

• オフマニフォールド・ドリフト（分布外への逸脱）: ステアリングされたノイズが、デコーダの訓練に使用された標準ガウス分布のサポートから逸脱し、デコーディングを不安定にする可能性があります。
• モード崩壊: 学習された潜在方策が、バックボーンのマルチモーダルな振る舞いを狭いサブセットのモードへと崩壊させてしまう可能性があります。

2. 手法: Lagrangian Perturbation Diffusion Steering (LP-DS)

著者らは、デコーダの重みを変更したり事前分布を完全に置き換えたりするのではなく、状態に条件付けられた潜在ノイズ空間内の残差を学習することで、凍結された生成方策を軽量に適合させるフレームワークである LP-DS を提案します。

コアメカニズム

LP-DSは、凍結された生成バックボーン $\Phi$ を、潜在ノイズ $w$ と状態 $s$ をアクションへと写像するブラックボックス・デコーダとして扱います。ノイズを事前分布 $p_0$（例：$\mathcal{N}(0, I)$）から直接サンプリングする代わりに、LP-DSは学習可能な摂動ネットワーク $\Delta_\theta(s)$ を導入します：
$$w = \epsilon + \Delta_\theta(s), \quad \text{ここで } \epsilon \sim \mathcal{N}(0, I)$$
$\Delta_\theta(s) = 0$ のとき、この方策は元の行動クローニングの分布を正確に再現します。

ラグランジュ緩和による制約付き最適化

摂動がデコーダが訓練されたマニフォールドからノイズを押し出すのを防ぐため、LP-DSは学習問題を制約付き最適化として定式化します。目標は、摂動の大きさに信頼領域制約を課した上で、生成されたアクションの期待Q値を最大化することです。

著者らは、ステアリングされたノイズ分布と事前分布の間の期待カルバック・ライブラー（KL）ダイバージェンスを、シフトの二乗ノルムを用いて近似します：
$$\mathbb{E}_{s \sim D}[\|\Delta_\theta(s)\|^2_2] \leq \delta$$
ここで $\delta$ は、許容される偏差を制御するハイパーパラメータです。

これは ラグランジュ緩和 アプローチを用いて解かれます。制約を動的に強制するために、双対変数（ラグランジュ乗数）$\alpha$ が導入されます：
$$\mathcal{L}(\theta, \alpha) = \mathbb{E}[Q^W(s, w) - \alpha(\|\Delta_\theta(s)\|^2_2 - \delta)]$$
アルゴリズムは以下を交互に実行します：

1. アクター更新: $\mathcal{L}$ に関して $\theta$ について最大化（勾配上昇法）。
2. 双対更新: $\alpha \geq 0$ を強制するために、$\alpha$ に関して $\mathcal{L}$ を最大化（投影勾配上昇法）。摂動が信頼領域に違反した場合、$\alpha$ が増加し、アクターにペナルティを与えて事前分布により密着するように強制します。

アルゴリズムの構成要素

訓練手順（アルゴリズム1）には以下が含まれます：

• データ収集: ステアリングされた方策 $w = \epsilon + \Delta_\theta(s)$ を用いてアクションをサンプリングする。
• クリティック更新: テンポラル・ディファレンス（TD）学習によるアクション値クリティック $Q^A$ の更新、およびベースとなるノイズ分布上の $Q^A$ からの蒸留による潜在値クリティック $Q^W$ の更新。
• 原始・双対更新: 報酬最大化と事前分布の保持のバランスをとるために、摂動ネットワーク $\theta$ と乗数 $\alpha$ を更新する。

3. 主な貢献

本論文は、以下の4つの主要な貢献を述べています：

1. LP-DS フレームワーク: 状態条件付きの残差ノイズ制御を通じて、凍結された生成方策を改善する軽量な適応手法であり、フルデコーダのファインチューニングに伴う不安定性とコストを回避します。
2. ラグランジュ信頼領域目的関数: 摂動の大きさを動的に制約する新しい目的関数を提供し、報酬最大化と事前学習された潜在事前分布の保持との間のバランスを制御する仕組みを提供します。
3. モード崩壊の軽減: LP-DSが、従来の事前分布置き換え型ノイズ空間ステアリング手法（DSRLなど）と比較して、より高いアクション空間の多様性を維持できることを実証しました。
4. 汎用性: コンパクトな拡散方策を超えて、フロー・マッチングのバックボーン、大規模な視覚言語行動モデル（$\pi_0$）、および物理的なFrankaロボット実験への適用可能性を検証しました。

4. 実験結果

著者らは、RoboMimic（操作）、OpenAI Gym（移動）、Adroit（巧緻な操作）のベンチマークにおいて、LP-DSをDSRL、DPPO（方策勾配ファインチューニング）、およびオフライン・トゥ・オンライン拡散Q学習のベースラインと比較評価しました。

• パフォーマンス: LP-DSは一貫して成功率とエピソード報酬を向上させました。特に WALKER2D 環境において、最強のベースラインに対して約 25% のリターン向上 を達成しました。
• サンプル効率: フルデコーダではなく軽量な摂動モジュールのみを更新することにより、LP-DSは優れたサンプル効率とより安定した学習ダイナミクスを示しました。
• 多様性の保持:
  • 玩具的なマルチゴール・ナビゲーションタスク において、LP-DSは信頼領域の境界 $\delta$ を通じて、パフォーマンスと多様性のトレードオフを明示的に制御できました。小さな $\delta$ はマルチモーダルな振る舞いを保持し、大きな $\delta$ は特化を可能にします。対照的に、DSRLは即座に単一のモードへと崩壊しました。
  • AVOIDING（障害物回避）環境 において、LP-DSは小さな $\delta$ で広範な実行可能な障害物回避ルート（軌跡レベルのマルチモーダリティ）を保持しましたが、DSRLは単一の支配的な経路へと崩壊しました。
  • エントロピー分析: LP-DSはオンライン適応中、DSRLよりも一貫して高いアクション空間エントロピーを維持しており、これはラグランジュ制約が早期の崩壊を効果的に抑制しつつ、性能を向上させていることを示しています。
• アブレーション研究: ラグランジュ・ペナルティを除去すると、学習が不安定になり成功率が低下しました。ラグランジュ・ペナルティとノイズ境界の両方を除去すると、極めて不安定な挙動と低い成功率となり、信頼領域制約が分布外へのステアリングを防ぐ主要なメカニズムであることが確認されました。
• 感度: 本手法は、信頼領域のターゲット値 $\delta$ の特定の値に対して頑健であることが判明し、0.1を超える値が概して強力なパフォーマンスをもたらしました。

5. 重要性と主張

本論文は、LP-DSが生成方策のフルファインチューニングに代わる、安定かつ効率的な代替案であることを主張しています。凍結されたバックボーンを固定された事前分布として扱い、有界な残差シフトのみを学習することで、本手法は以下のを実現します：

• オフマニフォールド・ドリフトの防止: 潜在クエリをデコーダの有効な動作範囲内に留めることで、制約のないノイズ空間ステアリングで見られる不安定な挙動を回避します。
• マルチモーダリティの保持: 事前学習された事前分布を置き換える際に観察されるモード崩壊を軽減し、凍結されたバックボーンが持つ振る舞いの多様性を維持します。
• アーキテクチャを横断する汎用性: このアプローチは拡散モデルに限定されず、フロー・マッチングモデルや大規模な視覚言語行動モデルを含む、潜在からアクションへの写像を持つあらゆる生成方策に適用可能です。
• 物理実装への展開: シミュレーションベースの潜在適応を用いた物理的なFrankaロボット実験によって検証され、シミュレーション上のベンチマークを超えた適用可能性が示されました。

著者らは、LP-DSを、事前学習されたマルチモーダルな構造の維持という必要性と、タスク固有の最適化の必要性とを両立させる、高容量生成方策の「オンライン改善」のための実用的なソリューションとして位置付けています。