【DL輪読会】Simplifying Latent Dynamics with Softly State-Invariant World Models

Simplifying Latent Dynamics with Softly State-Invariant World Models Tomoshi Iiyama, Matsuo Lab 1

書誌情報 Simplifying Latent Dynamics with Softly State-Invariant World Models ● Tankred Saanum, Peter Dayan, Eric Schulz ● Max Planck Institute for Biological Cybernetics （ドイツ） ● NeurIPS 2024 ● 論文 : https://arxiv.org/abs/2401.17835 実装 : 非公開

• https://arxiv.org/abs/2401.17835

書誌情報 Simplifying Latent Dynamics with Softly State-Invariant World Models 概要 ● 世界モデルにおける潜在空間のダイナミクス学習を改善する正規化手法 PLSM を提案 ● 行動が環境に与える影響 （潜在状態の変化） をより予測しやすくなるように情報ボトルネックを導入 ● 既存手法と組み合わせることで、RLタスクでのパフォーマンスや汎化性能、長期予測性能が向上

背景

世界モデル ● 環境の遷移を予測する生成モデル ● 人間の小脳に形成されているといわれている内部モデルにヒントを得ている ● David Ha の “World Models” [Ha+ 18] が火付け役となり、研究が進められている ● 強化学習と組み合わせた “Dreamer” [Hafner+ 19] などの手法が有名 もし左に動かしたら どうなる...？ 潜在状態 自分の行動で条件付けた 未来を予測する 行動 「パドルを左に動かす」 実際の観測 現在時刻 次時刻 （想像）

行動が環境に与える影響 通常、エージェントのとる行動が環境に及ぼす効果は予測しやすい ● 「アクセルを踏む」 車が加速する 「ジョイスティックを右に倒す」 キャラクターが右に移動する 一方で、例外もある ● 「アクセルを踏む」 「ジョイスティックを右に倒す」 エンジンが切れている状態では 壁がある場合は 車は加速しない キャラクターが移動しない ● 同じ行動をとっても、状態によって環境に与える影響が大きく異なる場合がある ● 現状の世界モデルでは、エージェントの「行動の効果」を予測しやすい形で体系的に扱えていない

行動が環境に与える影響 ● モデルが 「潜在状態の変化量」 を予測することで次時刻の潜在状態を決定するケースを考える 次状態 ● 現在状態 ある行動をとった時に起こる 状態変化の予測 行動の効果 （＝状態変化） が現状の世界モデル内で必ずしも 簡潔・単純に 表現されていない 似たような状態でも、同じ行動を実行した際に、大きく異なる状態変化を生成する場合がある ○ ○ これはエンコーダーが観測からどのような潜在空間を学習するかに大きく依存する → 行動の変化が一貫して表現される形に学習されるとは限らない 次状態を直接予測する モデルではない

行動が環境に与える影響 ● エンコーダーによって学習される潜在表現 ○ 平面上を移動する物体の例 ○ 物体 （赤） の x, y 座標を2次元の潜在状態で表現 ○ 状態変化は、 8方向に動く + 壁に当たって動けない ＝ 9通りの可能性がある 通常 状態変化を効率的に 表現するよう学習した場合 それぞれの次元の状態変化の分布

提案⼿法

通常の潜在ダイナミクス ● 状態（観測） をエンコーダーでコンパクトな潜在状態へマッピング ● 選択した行動をもとに引き起こされる状態変化を予測 行動の表記を省略 ● 潜在状態の遷移

簡素な潜在ダイナミクス ● ある行動が与えられた時の状態変化の確率分布 （潜在状態に関して周辺化） ● ある行動が、環境の遷移を予測可能な情報を持っている場合を考える ● この場合の のエントロピーは低くなる ○ ＝ 行動さえわかれば次の状態変化をほぼ予測できる ○ 分布が尖り不確実性（エントロピー）が低くなる ● 潜在状態を知っても、行動が与える効果についてほとんど情報を与えない ● できるだけ行動のみから世界の変化を正しく予測するように学習できないか？

簡素な潜在ダイナミクス ● ● ● 相互情報量で関係を考える 潜在状態と状態変化 行動のみから 行動＋潜在状態から の間の相互情報量 得られる状態変化の情報 得られる状態変化の情報 相互情報量が 0 の時 （極端な場合） ○ ダイナミクス（状態変化） はすべて行動のみに依存する ○ 潜在状態は無駄な情報 この相互情報量を最小化するように潜在表現を学習したい

簡素な潜在ダイナミクス ● ● 状態変化の予測を以下のように変更 ○ クエリネットワーク ○ 潜在コード を経由して状態変化を予測 クエリネットワークは、潜在状態 からダイナミクス予測に必要な情報のみを絞り込んで取り出す ○ 状態変化に影響を与える情報のみを許可したい ○ 潜在状態は状態変化の予測に直接使われない ■ ● を導入: 状態行動のペアを一旦 「潜在コード」 へマッピング → もし必要な情報があるなら に含まれる に含めなければいけない の情報を最小限にするには？

ダイナミクスへの情報ボトルネック ● ● に状態変化予測に必要な最小限の情報だけを含めたい → L2ノルムのペナルティを課す この正則化により、先ほどの相互情報量が最小化される ○ ● に の情報ができるだけ含まれないようになる ○ この正則化が、平均0、分散一定の正規分布とのKL最小化と等価であることが示されている [Ghosh+ 2019] ○ 詳細は Appendix を参照 潜在状態に直接ボトルネックを適用するのではなく、ダイナミクスに対して適用している ○ 行動から完全にダイナミクスが予測できる場合は、 の情報は一切失われない ○ → 次状態の予測に関して直接制約をかけているわけではない ○ RSSM や VAE のKL項は潜在表現自体にボトルネックを適用することに相当 この部分にのみ制約

Parsimonious Latent Space Model (PLSM) ● この正則化により世界モデルを学習する手法: Parsimonious Latent Space Model ○ ● “Parsimonious” : 倹約する、単純な 他の世界モデル手法に提案の正則化項を導入して性能を検証

実験

実験 - モデルベースRL ● ● TD-MPC [Hansen+ 2022] ○ 近年高い性能を示している世界モデル ○ 自身の潜在表現を予測するため、Self Predictive Representation (SPR) モデル というカテゴリに分類される TD-MPC の世界モデル “TOLD” のロスに少し変更を加える ○ ダイナミクスの部分を変更 + 正則化項を追加

実験 - モデルベースRL ● TD-MPC + PLSM (提案手法) ● DeepMind Control Suite (DMC) で検証 ○ 連続値制御タスク ○ 5種類の環境を選択 ● エピソードのスコアで比較 ● 1つのタスクを除いて、通常のTD-MPC よりも高いスコアを達成

実験 - 視覚的外乱があるタスク ● Repo [Zhu+ 2023] ○ ● ● DreamerをベースにしたモデルベースRL手法 の代わりに を使って GRU を更新するように改修 Distracting Control Suite で検証 ○ DMC をベースにした環境 ○ 背景がランダムな動画に置き換えられている ○ 動画はエージェントの行動とは独立で、報酬とも関係がない

実験 - 視覚的外乱があるタスク ● Repo + PLSM ● 各エピソードのスコアで比較 ● 難しいタスクでベースラインを大きく改善 ● 行動の効果を一貫して表現することで、外乱のある環境でも汎化能力を改善可能

実験 - モデルフリーRL ● 学習した表現はモデルフリーRLでも有用か？ ● TD-MPC と同様に潜在表現を学習 → Atari タスクで検証 ● いくつかのタスクで人間のスコアを上回る結果に ● 正則化が悪影響を及ぼすタスクも存在 ○ 正則化の強度を調整することによって改善するかも

実験 - 将来予測性能 ● 世界モデルの長期予測性能を改善できるか検証 ● 予測用のデータセットを作成 ● ○ Atari ○ 2D/3D 物体の操作 ○ dSprite ○ MOVi-E （行動なし） ベースライン ○ CWM, C-SWM [Kipf+ 2023] ○ 対照学習を利用した Object-centric の世界モデル

実験 - 将来予測性能 ● それぞれの予測ホライズンにおける予測精度で評価 ● 多くのデータセットでより遠くの将来まで予測することができた ● Slot-based の手法でも改善がみられた

実験 - 長期予測の汎化性能 ● ● テスト時のデータを変更して汎化性能を検証 ○ 訓練時よりも物体の数が少ないデータ ○ ガウスノイズをのせた画像 ドメインシフトにより全般的に精度は低下するが、それでも分布外のデータに対して良い汎化性能を示した

まとめ ● 行動の効果を簡潔に表現することのできる世界モデルを提案 ○ 潜在状態をできるだけ用いずに状態遷移を予測する正則化を導入 ● 長期予測性能、汎化性能、プランニング性能において一貫した改善が確認された ● 行動が環境に対して決定論的な効果をもたらすケースを想定している ● ○ 行動の条件付けとは無関係な環境の振る舞いは無視される傾向にある ○ 下流タスクに関連する情報も無視されてしまう場合がある 正則化の強度をタスクごとに適切に調整する必要がある