EXPO: Stable Reinforcement Learning with Expressive Policies

以下は、論文「EXPO: Stable Reinforcement Learning with Expressive Policies」の詳細な技術的サマリーです。

1. 問題定義

本論文は、オフラインデータセットが与えられた条件下で、オンライン強化学習（RL）を用いて表現力のあるポリシークラス（拡散ポリシーやフローマッチングポリシーなど）を微調整する課題に取り組んでいます。

• 核心的な対立: 表現力のあるポリシーは、模倣学習（IL）を通じて複雑な多モーダルな行動分布を捉えることに優れていますが、RL による直接的な最適化は困難です。単純なガウスポリシーとは異なり、表現力のあるポリシーは長いデノイジングチェーンによってパラメータ化されています。これらのチェーンを通じて価値勾配を逆伝播させてポリシーパラメータを更新することは、計算コストが高く不安定であり、トレーニングの崩壊や収束の不良を招くことがよくあります。
• ギャップ: 既存のオンライン RL 手法は単純なガウスポリシー向けに設計されており、事前学習済みの表現力あるモデルの能力を活用できていません。逆に、表現力のあるポリシーの微調整を試みる手法は、安定性の問題に直面したり、デノイジングプロセスに対して複雑な修正を必要としたりすることが多いです。
• 目標: 表現力のあるポリシーのパラメータを価値関数に対して直接最適化することなく、オフラインデータを用いて事前学習済みの表現力あるポリシーを微調整できる、サンプル効率が高く安定したオンライン RL アルゴリズムを開発すること。

2. 手法：EXPO（Expressive Policy Optimization）

著者は、表現力のあるポリシーを直接的な価値最大化から切り離すフレームワークであるEXPOを提案します。表現力のあるポリシーを通じて勾配を逆伝播させる代わりに、EXPO は Q 値を最大化する**「オンザフライ（OTF）」ポリシー**を構築します。これは 2 つの異なるコンポーネントを使用します。

A. 2 成分ポリシーアーキテクチャ

1. ベース表現力ポリシー（$\pi_{base}$）:
   • 役割: 行動分布に基づいて初期の行動候補を生成します。
   • トレーニング: 安定した模倣学習（IL）目的（例：Denoising Diffusion Probabilistic Models）を用いてトレーニングされます。Q 値を最大化するように明示的にトレーニングされることはなく、その安定性と複雑なモードを捉える能力が維持されます。
2. 軽量編集ポリシー（$\pi_{edit}$）:
   • 役割: $\pi_{base}$ からサンプリングされた行動 $a$ を受け取り、修正残差 $\hat{a}$ を出力して、洗練された行動 $\tilde{a} = a + \hat{a}$ を生成する、小型のガウスポリシーです。
   • トレーニング: 標準的なRL ポリシー勾配（$Q(s, a + \hat{a})$ の最大化）とエントロピー正則化を用いてトレーニングされます。
   • 制約: 編集はハイパーパラメータ $\beta$ によってスケーリングされ、元の行動に近くなるように制約されます。これにより、編集ポリシーがグローバルな最適化問題ではなく局所的な最適化問題を解決することが保証され、効率的かつ安定したトレーニングが可能になります。

B. オンザフライ（OTF）ポリシー抽出

サンプリングおよび時間的差分（TD）バックアップに使用される実際のポリシーは、単一のニューラルネットワークではなく、動的な選択プロセスです。

1. サンプリング: 与えられた状態 $s$ に対して、システムは $\pi_{base}$ から $N$ 個の行動をサンプリングします。
2. 編集: サンプリングされた各行動は $\pi_{edit}$ によって洗練されます。
3. 選択: システムは、元の行動（$a_i$）と編集された行動（$\tilde{a}_i$）の両方について Q 値を評価します。
4. 最大化: 候補の中で最も高い Q 値を持つ行動が実行のために選択され、TD バックアップのターゲットとして使用されます。

C. データ制限領域におけるエントロピーバックアップ

オフラインデータセットが小さいシナリオでは、著者はエントロピーバックアップを導入します。表現力のあるポリシー（拡散など）は閉形式のエントロピーを持たないため、Q 値で重み付けされた $N$ 個のベース＋編集行動候補にわたってソフトなサンプリング分布を構築します。これにより、解析的なエントロピー計算を必要とすることなく、RL 損失にエントロピーボーナスを含めることが可能になり、探索を促進します。

3. 主要な貢献

1. 安定した価値最大化戦略: 論文は、表現力のあるデノイジングチェーンを通じた直接的な勾配伝播を回避する新たなアプローチを導入しています。価値最大化を軽量な編集ポリシーと選択メカニズムに委譲することで、EXPO は直接最適化が失敗する状況でも安定性を達成します。
2. サンプル効率: EXPO は、オンライン RL およびオフラインからオンラインへの微調整の両方の設定において、先行する最先端手法と比較してサンプル効率で 2〜3 倍の改善を達成します。
3. ポリシー非依存性: このフレームワークは、ベースポリシーの特定のパラメータ化（拡散、フローマッチングなど）に依存せず、事前学習済みの価値関数を必要としないため、任意の事前学習済みロボットポリシーに適用可能です。
4. オンザフライ TD バックアップ: 著者は、TD ターゲット（サンプリングだけでなく）に価値最大化行動を使用することが学習速度にとって重要であることを実証しました。これにより、表現力のあるポリシーの文脈において、SARSA 風更新と Q-learning 風更新の間のギャップが効果的に埋められています。

4. 実験結果

この手法は、4 つのドメイン（AntMaze、Adroit、Robomimic、MimicGen）にまたがる12 の困難な連続制御タスクで評価されました。これらはいずれもスパース報酬を特徴としています。

• オンライン RL（事前学習なし）: EXPO は、サンプル効率においてベースライン（RLPD、IDQL、QSM など）を大幅に上回り、オフラインデータセットを事前学習に活用しなくても、より速く収束することがよくありました。
• オフラインからオンラインへの微調整:
  • EXPO は、他の手法でよく見られるパフォーマンスの崩壊（例：DAC の崩壊、IDQL の停滞）なしに、事前学習済みポリシーの微調整に成功しました。
  • 糸を通すやブロックを積み上げるなどの微細な操作タスクにおいて、他の手法が初期の模倣学習のパフォーマンスを超えて改善することに苦労する中、高いパフォーマンスを維持しました。
• アブレーション研究:
  • OTF TD バックアップ: TD ターゲットに対する価値最大化選択を除去し（サンプリングのみを使用）、パフォーマンスが劇的に低下しました。これは、バックアップにおける Q 値の整合性の重要性を確認するものです。
  • 行動編集: 編集ポリシーを除去すると、特に狭い行動分布を超えた探索を必要とするタスクにおいて、パフォーマンスが低下しました。
  • オフラインデータの質: パフォーマンスはオフラインデータセットの質と強く相関していました。しかし、エントロピーバックアップ変種は、模倣学習のみでは成功率が 10% 未満であった非常に質の低いオフラインデータであっても、EXPO が成功することを可能にしました。

5. 意義

EXPO は、「表現力のあるポリシーの微調整」というボトルネックを解決することで、ロボティクスと RLにおける重要な前進を表しています。

• 実用的な影響: 人間のデータを模倣することに優れた大規模な事前学習済み生成モデルを、オンライン相互作用を通じて適応し改善する必要がある実世界環境へ展開可能にします。
• 理論的洞察: 複雑な生成チェーンを通じて RL 用に勾配を逆伝播させる必要性に疑問を呈し、「局所的な編集＋グローバルな選択」という戦略の方が安定性が高く効率的であることを示しています。
• 将来の方向性: この研究は、オンライン設定で任意の事前学習済み生成ポリシー（拡散、フローなど）を微調整する扉を開き、ロボティクスにおける基盤モデルの採用を加速させる可能性があります。

限界: この手法は、バッチ更新ごとに複数の行動をサンプリングする必要があり、計算コストがかかります。さらに、合理的な事前情報（オフラインデータセットまたは事前学習済みポリシー）が存在することを前提としており、事前知識なしにゼロから学習することを目的としたものではありません。