[論文レビュー] Efficient Diffusion Policies for Offline Reinforcement Learning

Bingyi Kang, Xiao Ma+3 more|arXiv (Cornell University)|May 31, 2023

Robotic Locomotion and Control被引用数 5

ひとこと要約

本論文は、トレーニング中に汚れた状態からのアクションを近似することで、コストの高い拡散サンプリングチェーンを回避する、効率的拡散方策（EDP）を提案する。EDPは、Diffusion-QLと比較して25倍高速なトレーニング（5時間 vs. 5日）を達成し、サンプル効率と一般化性能を向上させることで、すべてのD4RLドメインで新たなSOTA性能を達成する。値ベースと尤度ベースのRLアルゴリズムの両方をサポートする。

ABSTRACT

Offline reinforcement learning (RL) aims to learn optimal policies from offline datasets, where the parameterization of policies is crucial but often overlooked. Recently, Diffsuion-QL significantly boosts the performance of offline RL by representing a policy with a diffusion model, whose success relies on a parametrized Markov Chain with hundreds of steps for sampling. However, Diffusion-QL suffers from two critical limitations. 1) It is computationally inefficient to forward and backward through the whole Markov chain during training. 2) It is incompatible with maximum likelihood-based RL algorithms (e.g., policy gradient methods) as the likelihood of diffusion models is intractable. Therefore, we propose efficient diffusion policy (EDP) to overcome these two challenges. EDP approximately constructs actions from corrupted ones at training to avoid running the sampling chain. We conduct extensive experiments on the D4RL benchmark. The results show that EDP can reduce the diffusion policy training time from 5 days to 5 hours on gym-locomotion tasks. Moreover, we show that EDP is compatible with various offline RL algorithms (TD3, CRR, and IQL) and achieves new state-of-the-art on D4RL by large margins over previous methods. Our code is available at https://github.com/sail-sg/edp.

研究の動機と目的

• 拡散方策のトレーニングにおける計算非効率性を解消すること。これは、長大なマルコフチェーンを経由する完全な前向きおよび後向きパスを必要とするためである。
• 尤度ベースのRLアルゴリズム（例：IQL、CRR）と拡散方策の互換性を実現すること。これらは、取り扱い可能な尤度を必要とする。
• 複雑で多次元なオフラインデータセットにおいて、標準のガウス方策を上回る一般性、効率性、拡張性を持つ方策パラメータ化法を開発すること。
• オフラインRLベンチマークにおいて、既存手法よりも高速かつ効果的であることを示すこと。

提案手法

• EDPは、再パラメータライゼーションを介して、汚れたアクションからクリーンなアクションへのマッピングを学習するノイズ予測ネットワークを、ノイズ除去拡散確率的モデル（DDPM）を用いて学習する。
• アクション近似を導入する。完全な拡散チェーンを経由するのではなく、学習済みのノイズ予測ネットワークを用いて、汚れた入力からアクションを再構築することで、トレーニング計算量を顕著に削減する。
• EDPは、固定分散とアクション近似から導かれる平均を持つガウス分布を用いて方策尤度を近似する。これにより、IQL や CRR などの尤度ベースのアルゴリズムとの互換性が実現される。
• オルソゴナル微分方程式（ODE）ベースのサンプラー、DPM-Solverを統合することで、サンプリングに必要なモデル呼び出し回数を削減し、トレーニングと推論の両方を高速化する。
• 推論時に、効率と性能のトレードオフとして10個のアクションを用いるエネルギーに基づくアクション選択（EAS）を適用し、サンプル品質を向上させる。
• TD3、CRR、IQL などの複数のオフラインRLアルゴリズムと、gym-locomotion、AntMaze、Adroit、Kitchen などのドメインを対象に評価することで、広範な互換性と性能向上を示している。

実験結果

リサーチクエスチョン

• RQ1 拡散方策のトレーニング時間を、性能を損なわず著しく短縮できるか？
• RQ2 IQL や CRR などの尤度ベースのRLアルゴリズムと、拡散方策を互換性を持たせられるか？これらは取り扱い可能な尤度を必要とする。
• RQ3 提案されたアクション近似法は、完全な拡散サンプリングと比較して、サンプル品質や方策性能を保持または向上させられるか？
• RQ4 EDPは、多様なオフラインRLアルゴリズムとベンチマーク環境に一般化可能であり、SOTA性能を達成できるか？

主な発見

• EDPは、D4RLのgym-locomotionタスクにおいて、トレーニング時間を5日から5時間に短縮し、Diffusion-QLと比較して25倍の高速化を達成した。
• EDPは、D4RLの4つのドメインすべてで新たなSOTA性能を達成し、RAT指標で平均スコア85.5（gym-locomotion）、63.0（AntMaze）、73.4（Adroit）、43.8（Kitchen）を記録した。
• AntMaze、Adroit、Kitchenでは、EDP + IQLがすべてのベースラインを平均スコアで10点以上上回り、複雑で多次元なタスクにおける強力な性能向上を示した。
• EDPは、TD3、CRR、IQLのすべてのアルゴリズムと組み合わせて使用しても、性能を維持または向上させ、すべてのアルゴリズムと環境で標準のガウス方策を上回ることを示した。
• アブレーションスタディにより、OMS（最良）スコアはトレーニングの不安定性のため誤解を招く可能性があることが確認され、RAT（平均）がより信頼性のある評価指標であることが示された。
• DPM-Solverで15回のモデル呼び出しを実行すると、ほぼ最適な性能が得られ、効率性も良好である。EASにおけるアクション数の増加は、9つのタスクのうち8つで性能を単調に向上させた。