【DL輪読会】Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics

Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics Tomoshi Iiyama, Matsuo Lab

1 書誌情報 Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics 離散潜在ダイナミクスからの適応的な時間的抽象化による階層世界モデルの学習 Christian Gumbsch, Noor Sajid, Georg Martius, Martin V. Butz Tübingen大学 (ドイツ) , UCL イギリス) ICLR 2024 Spotlight https://openreview.net/forum?id=TjCDNssXKU ( 2

• https://openreview.net/forum?id=TjCDNssXKU

Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics 概要 階層的な世界モデル THICK を提案（低レベル＆ レベルの2層） 低レベルは正確で細かい予測を維持しつつ レベルは時間的に抽象化された解釈可能な表現を学習 階層的な予測をモデルベース強化学習などのプランニングに利 → 従来の世界モデル（1層）と 較して性能が向上した 用 高 3 比 高 1 書誌情報

2 背景 世界モデル World Model) 間の 脳に形成されているといわれている内部モデルにヒントを得た、環境の遷移を予測する David Haの “World Models” Ha+ 18 が 付け役となり、研究が進められている 状態 表現 状態 表現 遷移の予測 推論 モデル 現在時刻 次時刻 状態 表現 状態 表現 成 動 観測 4 生 火 ] [ ( 小 行 人 生 成モデル

2 背景 世界モデルとプランニング 良い状態表現の獲得 外界の情報を 分に集約した低次元の特徴量を、学習によって獲得可能 → 強化学習の状態表現を 間が決める必要がなくなる 将来予測による意思決定 現実世界で実際に 動しなくても、反実仮想を → 遠くの未来まで先読みしてから、 成可能 番良さそうな 動をとる サンプル効率の向上 現実世界で実際に 動しなくても、反実仮想を → 想像したデータを使って強化学習の 成可能 策を訓練する（イメトレ） 行 生 生 方 一 人 行 行 十 5

世界モデルの課題 間はさまざまな時間スケールを柔軟に切り替えることが可能 パーティーを主催する 招待状を書く ペンを持ち上げる 筋 指令を送る 脳内には感覚運動経験を通じて複雑な階層的表現が形成されている 既存の世界モデルは、事前に固定されたタイムケールでの予測しかできない → 期 (long-horizon) タスクにおけるプランニング能 → 階層構造の要素を取り に限界がある れたい 力 入 長 6 肉 人 2 背景

時間的抽象化 Temporal Abstraction) 時間 向に抽象化された レベルの表現の獲得 法はまだ確 されていない レベル 低レベル 時間 事前固定の階層的な時間スケールに縛られていけないということが重要 どのように適切な時間スケールを発 し、抽象化するべきなのか？ 立 方 見 高 ( 7 方 高 2 背景

3 法 提案 C RSSM THICK RSSM を改良した 低レベルの世界モデル 階層世界モデル レベルが低レベル（C RSSM） 時系列データの中から スパースに変化する 潜在因 （context）を学習 の変化を予測 適応的な時間スケールで動作 プランニング THICKの予測を どうプランニングに いるか Dreamer および PlaNet と組み合わせる 8 - 用 × 子 子 - 手 高 の潜在因 THICK

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSMの潜在状態の値は、時間とともに連続的に変化する 時間 階層世界モデルの確 定 定 まばらに変化する 定 時間 立 ( - 手 一 一 一 には、時間的にスパースに変化するような潜在状態 ”context” が必要 9

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM min KL min KL ( - 手 10

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM ( - 手 11

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM ( - 手 12

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM ( - 手 13

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM ( - 手 14

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM min KL ( - 手 15

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM min KL ( - 手 16

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM min KL ( - 手 17

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM min KL ( - 手 18

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM min KL min KL ( - 手 19

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM 潜在状態 決定論的状態 確率的状態 Prior 確率的状態 Posterior 遷移 ( - 手 20

法 3 C RSSM Context-specific Recurrent State Space Model) 潜在状態 コンテキスト 確率的状態 細かいPrior 決定論的状態 確率的状態 粗いPrior 遷移 細かいダイナミクス Posterior 粗いダイナミクス ( - 手 21

法 3 C RSSM Context-specific Recurrent State Space Model) min KL min KL ( - 手 22

法 3 C RSSM Context-specific Recurrent State Space Model) ( - 手 23

法 3 C RSSM Context-specific Recurrent State Space Model) ( - 手 24

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM 確率的状態について、2つの予測を う 決定的状態 から求める正確で細かい prior コンテクスト から求める粗い prior （ を いない） 行 ( 用 - 手 25

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM 確率的状態について、2つの予測を う 決定的状態 から求める正確で細かい prior コンテクスト から求める粗い prior （ を いない） 行 ( 用 - 手 26

法 3 C RSSM Context-specific Recurrent State Space Model) 従来のRSSM 確率的状態について、2つの予測を う 決定的状態 から求める正確で細かい prior コンテクスト から求める粗い prior （ を いない） min KL 行 ( 用 - 手 27

法 3 C RSSM Context-specific Recurrent State Space Model) min KL min KL ( - 手 28

法 3 C RSSM Context-specific Recurrent State Space Model) 観測変数 画像 報酬 割引率 予測ヘッドも2種類 細かい予測 粗い予測 29 用 ( - 力 手 入 意する

法 3 C RSSM Context-specific Recurrent State Space Model) 粗いダイナミクス 粗いPrior 粗い予測 これらの予測経路は連続的に変化する → を使わずに から独 している のみによるroll-out（将来予測）が可能となる 立 ( - 手 30

法 3 GateL0RD Gumbsch+ 21 RNNの潜在状態を NeurIPS 同 著者 定にしたい → 変化量ができるだけ少なくなるように学習させたい L -ノルムによる正則化によってこれを実現 新しいコンテキストの提案値 コンテキストを更新するゲート コンテキストの更新式 ] 一 一 [ 1 2 0 2 0 手 31

法 3 GateL0RD Gumbsch+ 21 NeurIPS 同 著者 コンテキストの各要素の変化量を抑えたい 変化量のL -ノルム（ ゼロ成分の数）をペナルティとする そのままだと微分できない 状態が変化する ＝ 更新ゲートが開いている → 開いている更新ゲートの個数をペナルティとすればよい : Heavisideステップ関数 ] 非 一 [ 1 2 0 0 2 手 少しでも開いているゲートの 個数を数え上げる 32

法 3 C RSSM の学習 予測＆再構成 KLダイバージェンス コンテキストのL -正則化 0 - 手 33

法 3 THICK Temporal Hierarchies from Invariant Context Kernels) 低レベルの世界モデル （C-RSSM） をもとに レベルの世界モデル を構築したい 低レベルを抽象化 分割する 低レベル レベルのタイムスケールはどう決定するべきか？ → C-RSSMの離散的（スパース）に変化する context レベルの適応的なタイムスケールとして のダイナミクスを いる 核となっている仮定 コンテキストの更新は、潜在因 の重要な変化と 致している （特徴量の劇的な変化 ＝ 系列の適切な分割点とみなす） 34 一 用 子 ・ ( 高 手 高 高 高 レベル

3 法 THICK Temporal Hierarchies from Invariant Context Kernels) コンテキストが変化する位置が、意味のある系列の区切りである この重要な状態（変化点）を レベル 予測 予測 低レベル 時間 高 ( 手 高 レベルの世界モデルが予測する 35

3 法 THICK Temporal Hierarchies from Invariant Context Kernels) コンテキストが変化する位置が、意味のある系列の区切りである この重要な状態（変化点）を レベル 予測 予測 低レベル 36 高 ( 手 高 レベルの世界モデルが予測する

3 法 THICK Temporal Hierarchies from Invariant Context Kernels) 次にコンテキストが変化する（ 任意の ）タイムステップを表す関数 を定義 を次のコンテキスト変化点の時刻へマッピング 現在より後の時刻 1 かつ コンテクストの更新が起こる 2 3 4 5 6 7 8 9 低レベル ( 手 37

3 THICK Temporal Hierarchies from Invariant Context Kernels) レベルは低レベルのcontext の変化を引き起こすための変数（1ステップ前）を予測 確率的状態 動 38 ( 手 行 高 法

3 THICK Temporal Hierarchies from Invariant Context Kernels) レベルは低レベルのcontext の変化を引き起こすための変数（1ステップ前）を予測 確率的状態 動 39 ( 手 行 高 法

法 3 THICK Temporal Hierarchies from Invariant Context Kernels) レベルは低レベルのcontext プランニングに いるために以下の2つも予測する 経過時間 報酬和 40 用 ( 手 高 の変化を引き起こすための変数（1ステップ前）を予測

法 3 THICK Temporal Hierarchies from Invariant Context Kernels) レベルの予測時の 環境＆エージェントが従う 異なる結果が じる 41 方 力 入 生 ( 手 高 策に依存して

3 法 THICK Temporal Hierarchies from Invariant Context Kernels) レベルの予測時の Skillやoptionのような 動”を考える 42 行 高 力 入 ( 手 高 レベルの”

3 法 THICK Temporal Hierarchies from Invariant Context Kernels) Prior One-hot のカテゴリカル分布 Posterior Posteriorではゴール条件付けのエンコーディングを う Action State 行 ( 手 43

3 THICK の学習 低レベルに対する予測 レベルの 動 44 行 手 高 法

4 実験 THICKは時間的抽象化を学習できるのか？ → 学習された レベルのモデルは、様々なシナリオで意味のある解釈可能な抽象化ができていた THICKの階層的な予測でモデルベース強化学習 (MBRL) の性能は改善するか？ → 通常のDreamerよりも long-horizonタスクにおいて い性能を出した THICKの階層的な予測でモデル予測制御 (MPC) の性能は改善するか？ → フラットな世界モデルを いるよりも long-horizonタスクにおいて い性能を出した 高 高 用 高 45

4 実験 タスク 46

コンテキストの変化 コンテキストの更新がシーンの変化点と 致している 鍵をみつけてドアを開ける時（KeyRoom）、ドアノブを握る時（Door）など レベルはcontextの変化の1ステップ前を予測している（抽象化されているのでぼやける） 部屋の正確なレイアウトを与えられなくても、「部屋を出る」という抽象的な 動を予測している 行 47 一 高 4 実験

4 実験 コンテキストの変化 48

4 実験 コンテキストの変化 49

レベルの 2つの異なる 動 レベルの 動をとった時の予測 ボールを掴む、ボールを押す、パッドを踏むといった違う相互作 がエンコードされていることがわかる 用 行 行 50 高 高 4 実験

4 実験 モデルベースRL 51

4 実験 モデルベースRL 52

4 実験 MPC 53

階層世界モデルの構築のためにC-RSSMとTHICKを提案 抽象的なコンテキストの遷移と正確な低レベルのダイナミクスの両 → 階層的な予測により 予測がフラット ＆ 策が階層的（Director）よりも THICKはプランニング がうまくいく 法の提案ではない そもそもDreamer等が解けないタスクは、THICKを 論 内では2層の を予測可能 期的タスクの学習を改善できる 策がフラット ＆ 予測が階層的（THICK）の いても解けない 法を提案したが、N層のレベルに再帰的に拡張することも可能とのこと 方 用 方 長 手 方 手 54 文 方 5 まとめ