【DL輪読会】AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners

DEEP LEARNING JP [DL Papers] AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners Makoto Kawano (@mkt_kwn), Matsuo Lab. http://deeplearning.jp/

• http://deeplearning.jp/
• http://twitter.com/mkt_kwn

書誌情報 • AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners ▪ Liang, Z., Mu, Y., Ding, M., Ni, F., Tomizuka, M., and Luo, P ▪ The University of Hong Kong, University of California, Berkeley, Tianjin University, Shanghai AI Laboratory 今回のメイン ▪ ICML2023(oral) • Planning with Diffusion for Flexible Behavior Synthesis ▪ Janner, M., Du, Y., Tenenbaum, J.B., and Levine, S. 前回少し触れたが かなり簡素だったので ▪ University of California, Berkeley, MIT ▪ ICML2022 2

学習済みモデルによるプランニング • 強化学習やデータ駆動型の意思決定において非常に重要 ▪ 教師あり学習に近い未知の環境ダイナミクスの近似など ▪ うまくいくことはかなりレアケース：学習モデルがうまくいかないことが多いため ✓ うまくいくためには，因果的過程に従う必要がある ✓ 賢くない学習モデルを使って計画立ててもうまくいくわけがない • 本研究：軌道最適化をうまく扱えるモデルを作ること ▪ モデルからのサンプリングとモデルからの計画を同一的に扱う ✓ 状態ダイナミクスと行動の分布は重要 ✓ １ステップずつの精度より全体(長期)的な精度が大事 ✓ 一方で未知のタスクを含めた複数タスクに対応するため，z報酬に対して独立になるべき 3

問題設定 学習による軌道最適化を用いた動作合成 離散時間におけるダイナミクスを持つ系𝑠𝑡 = 𝑓 𝑠𝑡 , 𝑎𝑡 を考える．タイムステップごとの報酬（もし くはコスト）𝑟(𝑠𝑡 , 𝑎𝑡 )に因数分解された目的関数𝒥の最大化(もしくは最小化)するような行動の系列 𝑎∗0:𝑇 を見つけることを軌道最適化と呼ぶ： ただし，𝑇は計画期間を表し，𝜏 = (𝑠0, 𝑎0 , 𝑠1 , 𝑎1 , … , 𝑠𝑇 , 𝑎𝑇 )を状態と行動を交互にまとめたものとし ，𝒥(𝜏)を軌跡𝜏の目的値とする． 本研究の目標：軌跡𝜏を生成できるようなモデルを設計すること 4

軌道最適化における最大の課題 • 環境ダイナミクス𝑓に関する知識が必要であること ▪ 多くの手法：学習によるダイナミクスモデルで近似して計画に利用する ✓ 真のモデルを利用することを前提にした計画アルゴリズムが多い 学習モデルで計画しても敵対的な軌道を見つけてしまう • モデリングと計画を同時に行う拡散モデル𝑝𝜃 (𝜏)を利用 ▪ ノイズ除去の際に摂動分布からのサンプリングで条件づける=>確率論の問題になる 「Diffuser」と呼ぶ ✓ ℎ 𝜏 ：理想の出力や報酬，観測などに関する情報を含めた関数 ✓ 𝑝𝜃 (𝜏)のもとで物理法則的に正しく，ℎ 𝜏 のもとで高い報酬となる軌道を見つける 5

軌道計画における生成モデルDiffuserの要件：時間方向の順序づけ • サンプリングと計画を同時に行うと普通ではない制約が生じる ▪ 自己回帰による時間方向の状態予測ができなくなる ▪ ゴールで条件づけられた推論𝑝 𝑠1 𝑠0 , 𝑠𝑇 )を考える ✓「次の状態𝑠1 が前の状態𝑠0 と未来𝑠𝑇 に依存している」 ダイナミクスの予測：因果的 意思決定と制御：反因果的 => 過去によって現在が決まる => 未来に向かって現在を決める 全てのタイムステップにおける計画を同時に行う必要がある 6

軌道計画における生成モデルDiffuserの要件：時間方向の局所性/軌道の表現 • 自己回帰的/マルコフ的ではないが，時間的局所性を持たせたい ▪ ある時刻の予測：直近の過去/未来の情報にもとづく ✓1回のノイズ除去では局所的な一貫性にもとづく ✓繰り返すことで大域的な一貫性を持つことが可能 • 状態予測の質と同等に行動の有効性も重要 ▪ 状態と行動を同時に予測するように設計 ✓状態表現に行動表現を連結： 計画期間方向 7

Diffuserのアーキテクチャ • 基本はU-Netの1d-Convバージョン GN Mish Conv1D x Conv1D t FC Layer ▪ 畳み込みのおかげで計画期間の長さは可変になる GN, Mish • 学習も[Ho+, 2020]を利用 ▪ 𝑖 ∼ 𝒰{1,2, … , 𝑁}：拡散方向のタイムステップ ▪ 𝜖 ∼ 𝒩(0, 𝐼)：ターゲットノイズ 8

「報酬」導入による Diffuserでの強化学習 • Control-as-inference[Levine, 2018]と同じように定式化 ▪ 𝒪𝑡 ：時刻𝑡における軌道の最適性を表す二値確率変数 • 摂動関数ℎ 𝜏 = 𝑝(𝒪1:𝑇 |𝜏)によって最適な軌道をサンプリング可能 ガウス分布で近似 9

Diffuserの学習とプランニング • 学習手順： 1. 学習データセットを用いて拡散モデル𝑝𝜃 (𝜏)を学習 2. 各𝜏 𝑖 における累積報酬和を予測するモデル𝒥𝜙 を学習 • プランニング手順： ▪ 𝒥_𝜙の勾配でガイド 10

一般的な方法との特徴的な違い１：学習による長期期間計画 • 単ステップ予測モデル：真の環境ダイナミクス𝑓の近似利用を想定 ▪ 特化して計画アルゴリズムと紐付いているわけではない • Diffuser：拡散モデルのアフォーダンスに紐付いた計画手順 ▪ 計画≒サンプリングのため，diffuserの予測性能≒良質な長期計画 報酬がスパースな長期期間の計画 • 狙い撃ち法：近眼的な失敗が起きがち • Diffuser：うまくいく 11

一般的な方法との特徴的な違い2：時間方向の合成 • 単ステップモデル：マルコフ性を利用することが多い ▪ 学習データ分布の遷移を一般化して分布外データに適用 • Diffuser：局所的一貫性の反復=大域的一貫性を持つ ▪ いい感じで複数の軌道を組み合わせられる • 狙い撃ち法：直線的な軌道しか計画できない • Diffuser：二つの軌道をうまく組み合わせて V字軌道を計画できる 12

一般的な方法との特徴的な違い3/4：可変な計画期間/タスクの合成 • 畳み込みを利用しているため，任意期間の計画が可能 ▪ 入力ノイズ𝜏 𝑁 ∼ 𝒩(0, 𝐼)の長さNによって決められる • Diffuser：環境のダイナミクスと行動の両方に関する情報を含む 一方で報酬関数とは独立 ▪ モデル𝑝𝜃 𝜏 は，将来ありうる行動を出力 ▪ 摂動関数ℎ(𝜏)(複数組み合わせもあり)によって誘導される 13

(ここからAdaptDiffuser)Diffuserの欠点 そもそも拡散モデルの学習がうまくいかないとだめ =>学習データの多様性によって制限されがち ※特に初見タスクは解けない RQ：拡散モデルの強力な生成能力を用いて， 自身を改善するような異種データを作ることができるのか？ 14

AdaptDiffuser 1. とにかく軌道生成 � (Initialized as noise) Diverse Task Generation Goal 1 Start State � Denoising Process Denoising U-Net � Goal 2 ▪ 様々なタスク報酬でガイド Goal 3 Goal Reward Gradient Guidance Reward Function Update Diffusion Model Data Pool Drop Accept Start State Goal Discriminator � 15

AdaptDiffuser 1. とにかく軌道生成 � (Initialized as noise) Diverse Task Generation Goal 1 Start State � Denoising Process Denoising U-Net � Goal 2 Goal 3 ▪ 様々なタスク報酬でガイド 2. 識別器𝒟で選択 Goal ▪ 1で生成した軌道が Reward Gradient Guidance Reward Function Update Diffusion Model Data Pool 基準を満たしているか判断 Drop Accept Start State Goal Discriminator � 16

AdaptDiffuser 1. とにかく軌道生成 � (Initialized as noise) Diverse Task Generation Goal 1 Start State � Denoising Process Denoising U-Net � Goal 2 Goal 3 ▪ 様々なタスク報酬でガイド 2. 識別器𝒟で選択 Goal ▪ 1で生成した軌道が Reward Gradient Guidance Reward Function Update Diffusion Model Data Pool Drop Accept 基準を満たしているか判断 3. 良質な合成データで学習 ▪ 拡散モデルをfine-tune Start State Goal Discriminator � 17

AdaptDiffuser 1. とにかく軌道生成 Diffusion Model Initialize ▪ 様々なタスク報酬でガイド 2. 識別器𝒟で選択 II. Model (re-)training Goal Point Diverse Goal Point I. Offline Trajectories (Single Goal) Reward Guidance ▪ 1で生成した軌道が 基準を満たしているか判断 3. 良質な合成データで学習 IV. Diverse synthetic data Selection by Discriminator III. Guided Trajectory Generation ▪ 拡散モデルをfine-tune 4. 所望の精度まで継続 18

報酬ガイドによる合成データ生成 • 異なるゴールと報酬関数を用いて様々なタスクを定義する ▪ 連続的報酬関数：MuJoCo ✓ 連続値から写像された確率による最適性を示す2値確率変数に準拠 Diffuserの𝑝(𝒪 = 1|𝜏)と同じ ▪ スパース報酬関数：迷路(Maze2D)のようにゴール地点のみ1 ✓ 勾配がディラックデルタ関数になるためガイドに利用できない サンプリングされた値を制約で置き換えてしまえばよい ▪ 二つの報酬関数の組み合わせ：迷路の途中で寄るところがある 途中の報酬点を𝑠𝑐 として，p-ノルムをガイド関数として追加 19

識別器𝒟：合成データが基準を満たしているか • 従来手法の報酬ℛ(𝑠, 𝑎)を軌道𝜏 = (𝑠, 𝑎)生成との同時予測 ▪ ダイナミクスの一貫性を解決できてなかった ✓遷移モデル𝑠 ′ = 𝒯(𝑠, 𝑎)による実際と予測された次の状態が逸脱 ✓しかもそれを実行可能な解(=軌道)と見なされていた AdaptDiffuserでは • 逆動力学モデルℐを用いた状態追跡制御：𝑎෤ 𝑡 = ℐ 𝑠𝑡 , 𝑠𝑡+1 を利用 ▪ 得られた実際に実行可能な行動𝑎෤𝑡 は，ロボットの動的制約を守る ▪ 行動𝑎ො𝑡 を使って次の状態を得る：𝑠ǁ 𝑡+1 = 𝒯(𝑠ǁ 𝑡 , 𝑎෤𝑡 ) ▪ 𝑠ǁ 𝑡+1 とあまり離れていない𝑠𝑡+1 を採用 20

実験１：Maze2D • 過去の軌道履歴からゴールまでの最短経路を見つけるタスク ▪ 3種類のレイアウト ▪ 報酬はゴールに到達した時のみ１与えられるスパース性 ※この経路はDiffuserのもの (a) Diffuser (b) AdaptDiffuser Hard Case 1 (Maze2D-Medium) (c) Diffuser (Failed) (d) AdaptDiffuser Hard Case 2 (Maze2D-Large) 21

実験2：MuJoCo • 異種データ生成による性能向上の検証 ※この可視化はDiffuserのもの ▪ 基本的には全てのタスクで性能が最も高い ▪ 特にMediumデータセットでは顕著 ✓データ生成によって良質な学習ができている ▪ 一方でMedium-Expertはすでに良質のため効果薄 22

実験３：初見タスクでの適応能力（Maze2D） • Maze2Dで途中で金貨を集めるというタスクを追加 ▪ 金貨がない時：Diffuser共に最短経路を生成してる ▪ 金貨追加時：共に経路を変更している ✓Diffuser：壁に衝突している ✓AdaptDiffuser：衝突なし，かつより滑らかな経路 (a) Diffuser (Collision) (b) AdaptDiffuser Original Setting (c) Diffuser (Collision) (d) AdaptDiffuser (No Collision) With Gold Coin on (4, 2) 23

実験３：初見タスクでの適応能力（KUKA） • Stackingタスクの学習データのみ用意 • Pick and Placeタスクへの適応能力を評価 ▪ セットアップ１：4つのブロックが全て床の上にある状態スタート ▪ セットアップ２：ランダムに積み上がった状態でスタート 24

（Diffuser）෥ 𝒑𝜽 𝝉 ∝ 𝒑𝜽 𝝉 𝒉(𝝉)の独立性 • 条件付けなし • 条件付けあり 25

Ablation Study • Fine-tuningの反復回数 ▪ 1回 < 2回のfine-tuning ▪ Medium-Expertは既に良質すぎるため， 効果がそんなにない • 良質なデータ量 ▪ 50%のデータ量x25%の学習回数の時， Diffuserの100%データ量x10%学習回数と同等 • 学習時間 ▪ (一応主張としては) データ生成+fine-tuning<<学習時間 26

(付録として)ChatGPTによるデータ生成 • データ生成時にLLMを使ったらうまくいくのでは？ ▪ 迷路のレイアウトも増やせば良質なデータセットが増える ▪ 手作業で作るのは大変 ▪ プロンプトエンジニアリングでできそう (a) Existing Maze (b) Generated Mazes with ChatGPT 27

発表まとめ • 強化学習へ拡散モデルを応用したICML論文を２本 ▪ 状態の予測と行動の計画を同時に行なってしまう ▪ 拡散モデルの性能の高さを利用して，Self-evolveなデータ拡張も可能 • 感想 ▪ 自己回帰型ではなく，全体での系を扱うのは賢い ✓世界モデルとの関係性はどうなのか ▪ 上海AIラボすごい ✓OpenMMLabとかもここが管理してる ✓自動運転界隈でもすごい成果出してる 28