拡散モデルをポリシーとして用いる強化学習の後訓練は、性能向上のためにますます重要になっている。しかし、既存の拡散ポリシー勾配法はしばしば不安定で、信頼性の高いポリシー改善を達成できない。arXivの最新論文「Diffusion Policy Optimization without Drifting Apart」は、この問題の根本原因を分析し、解決策を提案している。
論文の著者Haozhe Jiangら6人は、不安定性の核心を「二重ドリフト」現象にあると指摘する。変分代理目的を最適化する際、証拠下限(ELBO)が真の対数尤度から乖離し、その結果、計算される代理ポリシー勾配が期待リターンの真のポリシー勾配とずれてしまう。勾配方向が不正確になるため、ポリシー更新が誤った方向に進み、最終的に性能向上が妨げられる。
この問題を解決するため、研究チームはDiPODと呼ばれる拡散ポリシー最適化フレームワークを提案する。このフレームワークの中心的な考え方は、訓練全体を通じてタイトバウンド挙動を維持することであり、具体的には自己蒸留とポリシー改善勾配更新を交互に実行する。自己蒸留ステップは変分分布と真の分布の密接な関係を保ち、ELBOのドリフトを防ぐ。この考えに基づき、著者らはシンプルで実用的なアルゴリズムを導出した:各拡散ポリシー勾配更新に、オン・ポリシーのELBO正則化項を追加するのである。この正則化項は、更新後のポリシーが現在のポリシーから離れすぎないように制約し、代理勾配が常に真の勾配と一致することを保証する。
実験は、拡散言語モデルの後訓練と連続制御タスクにおける拡散ポリシーの2つの重要な応用シナリオで検証された。拡散言語モデル後訓練では、DiPODは訓練を大幅に安定させ、従来法よりも高い報酬を達成した。連続制御タスクでもDiPODは優れた性能を示し、その汎用性と有効性が証明された。
本研究は、拡散ポリシーの強化学習後訓練にシンプルかつ効果的な安定化手法を提供する。理論分析により二重ドリフトという根本問題を明らかにし、それに基づいて実用的な正則化技術を設計したDiPODは、拡散ポリシーの幅広い実応用展開を促進する可能性がある。論文のプロジェクトページでは詳細が公開されており、関連コードと実験設定も利用可能である。