技術的概要：適応的事前分布を用いた計画学習のための KL 正則化フレームワーク

1. 問題定義

モデルベース強化学習（MBRL）において、効果的な探索は依然として重要な課題であり、特にサンプル効率が決定的に重要となる高次元連続制御タスクにおいて顕著です。TD-MPC や BMPC などの最近のアプローチは、モデル予測経路積分（MPPI）計画のための提案分布として学習された方策を利用しています。しかし、方策のミスマッチという根本的な制限が生じます。すなわち、サンプリング方策（$\pi_{\theta_s}$）は、計画者の行動分布（$\pi_P$）とは独立して更新されることが多いのです。

この分離により分布のミスマッチが生じます。具体的には、ブートストラップ行動価値関数は MPPI 計画者によって生成された軌道（特定の分布を探索する）に基づいて訓練される一方、サンプリング方策は局所的な行動価値を最大化するように更新され、しばしば計画者の分布から乖離します。この不一致は、特に短い時間軸において信頼性の低い価値推定をもたらし、長期的な性能を低下させます。これらの分布を整合させるための以前の試み（BMPC における逆 KL 発散の最小化など）は、リプレイバッファに格納された「古くなった」計画者のサンプルに依存しているため、高分散に悩まされます。これらのサンプルは、一貫した現在の分布ではなく、時代遅れの計画方策の混合を表しているからです。

2. 手法：PO-MPC フレームワーク

著者らは、計画者の行動分布を方策最適化における事前分布として扱う統合フレームワークである方策最適化–モデル予測制御（PO-MPC）を提案します。PO-MPC は、サンプリング方策の学習ステップをKL 正則化強化学習の一例として定式化します。

中核的な定式化

目的は、適応的事前分布 $\pi_p$（MPPI 計画者から導出される）に近接しつつ、期待リターンを最大化するサンプリング方策 $\pi_{\theta_s}$ を学習することです。目的関数は以下の通りです：

$$J(\pi_{\theta_s}) = \mathbb{E} \left[ \sum_{t=0}^{T-1} \gamma^t \left( r(z_t, a_t) - \lambda \text{KL}[\pi_{\theta_s}(\cdot | z_t) \| \pi_p(\cdot | z_t)] \right) \right]$$

ここで：

• $z_t$ は潜在状態表現です。
• $\lambda$ は KL 正則化の強さを制御するハイパーパラメータです。
• $\pi_p$ は方策の事前分布です。

主要な構成要素

1. KL 正則化 RL： このフレームワークは、既存の手法をこの定式化の特殊なケースとして統合します。

   • $\lambda = 0$：TD-MPC2を再現します。ここでは、方策は行動価値の最大化とエントロピー正則化のみを通じて更新されます。
   • $\lambda \to \infty$：BMPC（逆 KL を用いた場合）を再現します。ここでは、方策は計画者の分布を模倣するように制約されます。
   • 中間の $\lambda$： 本論文では、リターン最大化と分布整合のバランスを取るために、中間的な値（例：$\lambda \in [0.1, 9]$）を探求します。

2. 適応的中間事前分布（$\pi_{\theta_p}$）：
   リプレイバッファから直接取得する古くなった計画者サンプルを使用することによる高分散に対処するため、PO-MPC は学習された中間事前分布を導入します。格納された計画者サンプルの生分布（時間変化するガウス混合分布となり得る）をそのままクローンするのではなく、現在の MPPI 誘導分布 $\pi_P$ を近似するニューラルネットワーク $\pi_{\theta_p}$ を訓練します。

   • この事前分布は計画者の出力から蒸留されるため、サンプリング方策の更新を時代遅れのサンプルの分散から保護します。
   • この事前分布は、逆 KL（$\text{KL}[\pi_{\theta_p} \| \pi_P]$ の最小化）または順 KL（$\text{KL}[\pi_P \| \pi_{\theta_p}]$ の最小化）のいずれかを用いて訓練できます。
     • 逆 KL： 方策を計画者の特定のモードに一致させるようにバイアスします（収束が速いが、探索が不足する可能性あり）。
     • 順 KL： 方策をすべてのサンプリング分布のサポートをカバーするようにバイアスします（探索が強化されるが、収束が遅くなる可能性あり）。

3. アルゴリズムループ：
   訓練は計画 $\to$ 推論 $\to$ 正則化のサイクルに従います。

   • 計画： 現在のサンプリング方策と価値関数を用いて、MPPI が軌道分布 $\pi_P$ を生成します。
   • 推論： 適応的事前分布 $\pi_{\theta_p}$ が $\pi_P$ を近似するように更新されます。
   • 正則化： サンプリング方策 $\pi_{\theta_s}$ は、Q 値の最大化と $\pi_{\theta_p}$ への KL ペナルティのバランスを取る KL 正則化目的関数を用いて更新されます。

3. 主要な貢献

• 統合： 本論文は、TD-MPC2 や BMPC といった多様な MPPI ベースの RL 手法を、単一の理論的フレームワーク（PO-MPC）の下で統合し、それらの設計選択と関係を明確にします。
• 適応的事前分布： リプレイバッファのサンプルを直接クローンする代わりに、学習された中間事前分布を導入します。これにより、勾配分散が減少し、訓練が安定化することが理論的および実証的に示されています。
• 柔軟な目的関数： このフレームワークにより、以前は研究されていなかった変種の探索が可能になります。具体的には、KL 強度 $\lambda$ の調整や、事前分布に対する順 KL/逆 KL 損失の選択です。
• 最先端の性能： 本手法は、既存のベースラインと比較して、高次元連続制御タスクにおいて優れたサンプル効率と最終性能を達成します。

4. 実験結果

著者らは、DeepMind Control Suite（Humanoid および Dog）とHumanoidBenchの 21 の困難な連続制御タスクにおいて PO-MPC を評価しました。

• 性能向上： 中間的な $\lambda$ 値を用いた PO-MPC は、TD-MPC2（$\lambda=0$）および BMPC（$\lambda \to \infty$）の両方を一貫して上回りました。
  • 高次元タスク（例：HumanoidBench）において、中間的な $\lambda$ 値（特に $\lambda=1$ または $\lambda=9$）は、最終リターンとサンプル効率において顕著な改善をもたらしました。
  • ベースラインが学習に失敗した一部の環境（例：HumanoidBench の「Balance Simple」、「Slide」、「Stair」）において、PO-MPC は正常に収束しました。
• 中間事前分布の有効性： 学習された中間事前分布（$\pi_{\theta_p}$）を使用することは、リプレイバッファからの生計画者サンプルを使用することと同等かそれ以上の性能を示し、更新分散の減少を確認しました。
• 損失関数の感度： 事前分布の訓練における KL 損失の選択は重要でした。探索が重要なタスク（例：「Stair」）では順 KL が有益であり、一方、精度が重要なタスク（例：「Balance Simple」）では逆 KL が収束を加速しました。
• アブレーション： 結果は、$\lambda$ が低すぎると（方策ミスマッチによるリターンの低下）、高すぎると（早期のモード崩壊）、性能が低下することを示しており、このトレードオフの調整の重要性を浮き彫りにしました。

5. 意義と主張

本論文は、PO-MPC が MPPI ベースのアプローチの原理的な統合を表すと主張しています。KL 正則化目的関数と適応的事前分布を通じてサンプリング方策と計画者を明示的に整合させることで、このフレームワークは以下の点を実現します。

1. 計画と方策学習の間のループを閉じ、価値関数が現在の計画者に関連する分布上で評価されるようにします。
2. リプレイバッファ内の時代遅れの計画サンプルによって導入される分散を軽減することで、安定性を向上させます。
3. 連続制御におけるモデルベース RL の最先端を進展させ、特定のタスク要件（探索対利用）に合わせて調整可能な柔軟なフレームワークを提供します。

著者らは、このフレームワークは強力である一方で、方策に対してガウス分布を仮定しており、価値学習のためにシミュレートされたすべての軌道を完全に活用していないと指摘しています。今後の課題として、ガウス分布を超えた方策分布への拡張と、$\lambda$ ハイパーパラメータの自動調整が提案されています。