Efficient Hypergradient Descent for Inverse Reinforcement Learning

技術要約：逆強化学習のための効率的なハイパーグラディエント降下法

問題提起
逆強化学習（IRL）は、エキスパートのデモンストレーションから報酬関数を復元することを目的としており、その報酬の下での最適方策が観測された行動を再現するように設計される。著者らは、最大尤度IRL（ML-IRL）に焦点を当てている。これは、バイレベル最適化問題として定式化できる。内側レベルでは、学習された報酬 $r_\phi$ の下でエントロピー正則化された方策を最適化し、外側レベルでは、誘導された方策とエキスパートデータの間の乖離（具体的にはKLダイバージェンス）を最小化する。

このバイレベル問題を暗黙的微分（implicit differentiation）を用いて解く際の主要な計算上のボトルネックは、ハイパーグラディエント（hypergradient）の計算を必要とすることである。これには、内側目的関数の逆ヘッセ行列ベクトル積（inverse-Hessian-vector product）が必要となる。ヘッセ行列を明示的に構築するには、二次的なメモリ（$O(d_\theta^2)$）を必要とし、一方で共役勾配法などの反復近似は、計算コストが高く、条件付け（conditioning）に敏感である。標準的な自然ハイパーグラディエント降下法（NHGD）は、ヘッセ行列をフィッシャー情報行列のサロゲートで置き換えることを提案しているが、標準的なNHGDは固定されたデータ分布を仮定している。ML-IRLでは、内側の目的関数は方策によって誘導される軌道分布に関する逆KLダイバージェンスを含むため、NHGDを直接適用することは非自明である。

手法
本論文は、内側問題の幾何学的構造を利用し、行列スケッチングを採用することで、ML-IRLのためのハイパーグラディエントを効率的に計算する手法を提案している。

1. フィッシャー・ヘッセ行列の等価性： 著者らは、厳密な内側最適性と実現可能性（方策クラスがエキスパート分布を完全にモデル化できるほど十分に豊かであること）の下で、内側目的関数のヘッセ行列が、割引軌道フィッシャー情報行列（$F_\theta$）に比例することを証明した。具体的には、$\nabla^2_\theta L_{inner} = \alpha F_\theta$ となる。これにより、ハイパーグラディエントに必要な逆ヘッセ行列ベクトル積を、逆フィッシャー行列ベクトル積で置き換えることが可能になる。
2. ハイパーグラディエントの推定： この手法は、ハイパーグラディエントの計算を以下の3つのコンポーネントに分解する：
   • 方策パラメータに関する外側勾配。
   • フィッシャー行列と外側勾配を含む減衰線形システムの解。
   • システムの解と報酬勾配を含む混合微分項。
3. ストリーミング・スペクトル・スケッチング（SCFD）： 密なフィッシャー行列を保持するための $O(d_\theta^2)$ のメモリコストを回避するため、著者らは、Streaming Spectral Compensation Frequent Directions (SCFD) スケッチを用いて、必要な逆フィッシャー行列ベクトル積を近似する。フィッシャー行列を実体化する代わりに、このアルゴリズムは、エージェントの軌道から重み付きの方策スコアベクトルを、サイズ $m$ のコンパクトなスケッチへとストリーミングする。これにより、ストレージ複雑度は $O(md_\theta)$ に削減される。その後、このスケッチを使用して、減衰線形システムを近似的に解く。

主な貢献

• 理論的導出： 著者らは、ML-IRLの暗黙的ハイパーグラディエントのためのサンプルベースの推定量を導出し、実現可能な内側最適において、内側ヘッセ行列が軌道フィッシャー情報行列に比例することを証明した。これは、ML-IRLにおける暗黙的ハイパーグラディエント計算にこの恒等式を適用した最初の事例として提示されている。
• スケーラブルなソルバー： 著者らは、方策スコアベクトルをストリーミングするSCFDベースのソルバーを提案している。このアプローチは、スコア行列や密なフィッシャー行列を構築することを回避し、メモリ要件を大幅に削減する。
• 実証的検証： 本手法は、離散（CartPole）および連続（LQR）制御環境において、第一次確率的バイレベル・ベースライン（シングルループML-IRL）および明示的フィッシャー・ソルバーと比較して評価されている。

結果
実験は、同一のアーキテクチャとハイパーパラメータ調整を行い、一致させた計算予算（24時間のトレーニング制限）の下で行われた。

• 方策の性能： フィッシャーベースの手法は、競争力のある方策品質を達成した。CartPoleでは、すべての手法がエキスパートレベルのリターンに達した。LQRでは、シングルループML-IRLが最高のポリシー負の対数尤度（NLL）および環境リターンを達成したが、「スケッチを用いたフィッシャー法」との差は僅かであった。
• 報酬のランキング： フィッシャーベースの手法は、強力な報酬ランキング品質を示した。LQRにおいて、「スケッチを用いたフィッシャー法」は、学習された軌道リターンと真の報酬との間のランク相関（RankCorrelation）において最高値（0.972）を達成し、シングルループのベースラインおよび明示的フィッシャー・ソルバーの両方を上回った。
• 効率性：
  • メモリ： 高次元のLQR設定において、スケッチングは最も大きな恩恵をもたらし、明示的フィッシャー・ソルバーと比較してピークメモリ使用量を最大1.31倍削減した。
  • 速度： CartPoleでは、スケッチングにより1.29倍の高速化を実現した。LQRでは、中程度のスケッチサイズ（例：$m=32$）は、明示的ソルバーに対してわずかな高速化（1.04倍）を提供したが、より大きなスケッチ（例：$m=256$）は、最終的な損失値は低いものの、スケッチングプロセスのオーバーヘッドにより、明示的ソルバーよりも遅くなった。
  • 安定性： 著者らは、スケッチされたソルバーが明示的な定式化よりも大幅に弱い減衰パラメータ（$\lambda$）の下でも安定していることを見出した。これは、スケッチングプロセスが暗黙的な正則化として機能している可能性を示唆している。

意義と主張
本論文は、逆KL最適化の文脈における内側ヘッセ行列とフィッシャー情報行列の間の直接的な関連性を確立することにより、バイレベルIRLのスケーラビリティのボトルネックに対処することを主張している。この理論的知見をストリーミング行列スケッチングと組み合わせることで、明示的なヘッセ行列やフィッシャー行列の構築に伴う二次的なメモリコストなしに、暗黙的ハイパーグラディエントを計算できることを著者らは示している。

著者らは、フィッシャーベースの手法が最終的な方策品質の面ではシングルループML-IRLと同等の競争力を持つ一方で、高次元の問題における報酬ランキングの性能と計算効率（特にメモリ削減）において明確な利点を提供すると控えめに結論付けている。本研究は、適度なスケッチサイズが、最適化の質、メモリ使用量、および実行時間の最適なトレードオフを提供することを示唆している。