yamatsuta_jsai_slide

潜在拡散モデルを用いた 世界モデルの提案 山蔦 栄太郎（大阪大学大学院工学研究科） 内山 史也 （東京大学） 関戸 麗矢 （北海道科学大学） 川原 雄登 （サンタモニカカレッジ） 鈴木 雅大 （東京大学大学院工学系研究科） 松尾 豊 （東京大学大学院工学系研究科）

概要 • World Models [Ha+, 2018]をHalf Cheaterで実行 • World Modelsの視覚処理に潜在拡散モデルを追加 • 制約条件下で既存のβ-VAEよりも高い報酬を獲得する ことに成功

背景・目的 • Making the World Differentiable [Schmidhuber 90] 環境を内部モデルとして表現し、環境の変化を予測しながら行動を決定 世界モデルの原型となる研究 • World Models[Ha+ 18] Schmidhuberのアイデアをもとに、 1. 観測データの圧縮 2. 将来観測するであろうデータの予測 3. 上記二つからより良い行動を選択する ことで観測データを内部表現として保有している 背景・目的

背景・目的 • Making the World Differentiable [Schmidhuber 90] 環境を内部モデルとして表現し、環境の変化を予測しなが ら行動を決定 世界モデルの原型となる研究 • World Models[Ha+ 18] Schmidhuberのアイデアをもとに、 1. 観測データの圧縮 2. 将来観測するであろうデータの予測 3. 上記二つからより良い行動を選択する ことで観測データを内部表現として保有している 背景・目的

背景・目的 背景・目的 • 世界モデル(World Models) 観測情報を圧縮するときのサブモデルが独立で、他の サブモデルにとって重要な部分が何か参照できない • 拡散モデル テキストなど様々なモダリティでの条 件付けを行える 拡散モデルで観測情報を条件付きで再構成することで、圧縮された観測情報を他の条 件(直前にとった行動など)を踏まえてより洗練させることができないか？

背景・目的 • 世界モデル(World Models) 観測情報は一つの独立なサブモデルによって圧縮され、 他のサブモデルにとって重要な部分は何か参照できない 時々タスクを解く上で関係ない情報がエンコードされたり重 要な部分の情報落ちが起こる 圧縮され欠損がある観測情報を他の条件(直前に とった行動など)を踏まえて補完できないか？ • 拡散モデル 背景・目的 圧縮 再構成 カーレーシングタスクにおける極端な例 条件 画像に対する高いサンプリング性能、元のデータに一部 欠損があっても条件をもとに修復ができる 修復

背景・目的 • 世界モデル(World Models) 観測情報は一つの独立なサブモデルによって圧縮され、 他のサブモデルにとって重要な部分は何か参照できない 時々タスクを解く上で関係ない情報がエンコードされたり重 要な部分の情報落ちが起こる 圧縮され欠損がある観測情報を他の条件(直前に とった行動など)を踏まえて補完することを提案 背景・目的

背景・目的 圧縮され欠損がある観測情報を他の条件(直前に とった行動など)を踏まえて補完できないか？ • 拡散モデル 画像に対する高いサンプリング性能、元のデータに一部 欠損があっても条件をもとに修復ができる 拡散モデルによる条件付け生成で潜在変数の欠損 を修復することを提案 背景・目的

World Models[Ha+ 18] • 環境を内部モデルで表現しそれに応じた行動を 取らせる手法 • Vision(VAE)とMemory(LSTM)で観測（画像） から環境の将来の表現を予測 • 一層の線形変換層で次の行動を決定 関連研究

Vision • Visionで観測情報を圧縮し潜在変数zを得る • zの出力以外で、他のサブモデルとは一切関係しない • β-VAE [Higgins+ 17]ではKLダイバージェンスに定数を掛け disentanglement(潜在変数の各要素の独立性)を調整している 関連研究

Memory 関連研究 • Visionで圧縮した観測情報を入力 • 次のステップにおけるzを予測 • MDN-RNNやLSTMが一般に使用されている

Controller 関連研究 • VAEの潜在変数とLSTMの隠れ状態が入力 • 一層の線形変換で行動を決定 • CMA-ESで最適化 • ランダムにサンプリングし上位サンプルの分布に合 うようにパラメータを調節

関連研究 DIffusion Model ノイズ • 画像などの多次元データにノイズを段階的に付加（拡 … 散過程） … • ノイズを除去する逆拡散過程を学習(DDPM) … • テキスト分散表現など条件付けでの生成も可能 DM

DIffusion Model • 画像などの多次元データにノイズを段階的に付加（拡 散過程） 関連研究 ノイズ … • ノイズを除去する逆拡散過程を学習(DDPM) … • 高い画像生成品質 … • 欠損した画像データをテキスト分散表現などの条件付 けのもと修復する用途にも使われる DM

DIffusion Model • 潜在拡散モデル(Latent Diffusion Model) [Rombach+ 22][Ramesh 22] 関連研究 ノイズ … • 変分オートエンコーダ(VAE)で特徴量をより低次元 の潜在変数にする … • 拡散・逆拡散をVAEの潜在変数上で行う … • 低次元で効率的にサンプリング可能 LDM VAE

手法概要 World Modelsの課題： VAEで圧縮した潜在変数においてタスクを解くにあたっ て不要な情報がエンコードされたり必要な情報が情報 落ちしている可能性 -> 潜在変数に他の条件を考慮することで修正 • VAEの潜在変数z’を潜在拡散モデルで学習 • hとActionの条件付けのもとz’をサンプリング • サンプリング結果をControllerに入力 提案手法

制約づけ 提案手法 Action 潜在拡散モデルのサンプリング時間が数秒程度で、 毎フレームサンプリングしていると訓練に時間がか かる（1エピソードに数秒 * フレーム数） ・CMA-ESのサンプル数を32->4に減少 ・LDMの更新制限 Full Sight: 毎フレームごとの入力でz’を更新 Limited SIght: 100枚に1フレームだけz’を更新、Controllerへの入 力に100フレーム同じものを使用し続ける

実験環境 実験 訓練環境: OpenGymHalfCheetah forward_reward(前進具合)から ctrl_lost(動作が大げさな時のペナルティ) を引いたものが報酬 モデル： Vision: betaVAE Memory: LSTM Controller: 線形層 観測フレーム: 64 * 64 LDM: 100ステップ NVIDIA A10G環境下で12時間以内に全ての訓練 が終了

実験

報酬の違い 平均報酬 実験 最大報酬 LDM VA E LDM VA E 最適化ステップ数

報酬の違い 平均報酬 実験 最大報酬 LDM VA E LDM VA E 最適化ステップ数 VAEと比べてLDMの報酬は高くかつ正値である -> 動作の大げさ具合によるペナルティが小さく、 また前進による報酬が得られていることが確認できる

実験 実際の挙動 LDM VAE LDMの場合動作が小さく、LDMの更新ステップ(100回に1回)の時に大きく動く

考察 • VAEで報酬が正値にならなかった CMA-ESにおけるサンプリング数を激減させたからと考えられる • 潜在拡散モデルだと制約条件下において比較的高い報酬を獲得できる(ペナルティが小さい) 課題 • 報酬が最大でも200程度だが、他手法が数千までの報酬獲得に成功している ->改善の余地が十分にある LDMの更新頻度を制限しているので、そもそもControllerへの入力の変化が小さすぎる？ パラメータの最適化や条件付ける特徴量の変更 • LDMを訓練するにあたって、学習済みVAEの存在が前提 タスクに重要な部分(HalfCheetahにおける足など)を重点的に再現できているかの検証

考察 • VAEで報酬が正値にならなかった CMA-ESにおけるサンプリング数を激減させたからと考えられる • 潜在拡散モデルだと制約条件下において比較的高い報酬を獲得できる(ペナルティが小さい) 課題 • 前進動作が小さかった LDMの更新頻度を制限しているので、Controllerに対する入力の変化が小さく、動くきっかけを 与えられていない可能性 Controllerへの入力条件の変更やパラメータの最適化、条件付ける特徴量の変更 • LDMを訓練するにあたって、学習済みVAEの存在が前提 潜在変数の再構成結果が、VAEとLDMでどのくらい違うものであるかの検証が必要

結論 • worldmodelsに拡散モデルを適用する手法を提案 • 拡散モデルのサンプリング時間を考慮したシーケンスを使用 • HalfCheetahにおいて比較的高い報酬を獲得 考察

参考文献 [Schmidhuber 90] Schmidhuber, Jiirgen.: Making the World Differentiable: On Using Self-Supervised Fully Recurrent Neural Networks for Dynamic Reinforcement Learning and Planning in Non-Stationary Environments. (1990). [Ramesh 22] Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M: Hierarchical Text-Conditional Image Generation with CLIP Latents, arXiv (2022), [Rombach 22] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B.: High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10684-10695 (2022), [Ha 22] Ha, D., and Schmidhuber, J: World Models, arXiv (2018), [Higgins 17] Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., ... and Lerchner, A.: betaVAE: Learning Basic Visual Concepts with a Constrained Variational Framework, ICLR (2017),

参考文献 [Ramesh 22] Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M: Hierarchical Text-Conditional Image Generation with CLIP Latents, arXiv (2022), [Wang 22] Wang, Z., Hunt, J. J., and Zhou, M.: Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning, arXiv (2022), [Rombach 22] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B.: High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10684-10695 (2022), [Schmidhuber 90] Schmidhuber, Jiirgen.: Making the World Differentiable: On Using Self-Supervised Fully Recurrent Neural Networks for Dynamic Reinforcement Learning and Planning in Non-Stationary Environments. (1990). [Ha 22] Ha, D., and Schmidhuber, J: World Models, arXiv (2018), [Pearce 23] Pearce, T., Rashid, T., Kanervisto, A., Bignell, D., Sun, M., Georgescu, R., ... and Devlin, S.: Imitating Human Behaviour with Diffusion Models, arXiv (2022), [Liang 23] Liang, Z., Mu, Y., Ding, M., Ni, F., Tomizuka, M., and Luo, P.: AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners, arXiv (2022), [Janner 22] Janner, M., Du, Y., Tenenbaum, J. B., and Levine, S.: Planning with Diffusion for Flexible Behavior Synthesis, PMLR (2022), [Kingma 14] Kingma, D. P., Mohamed, S., Jimenez Rezende, D., and Welling, M.: Semi-Supervised Learning with Deep Generative Models, NeurIPS (2014), [Higgins 17] Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., ... and Lerchner, A.: betaVAE: Learning Basic Visual Concepts with a Constrained Variational Framework, ICLR (2017), [van den Oord 18] Van Den Oord, A., and Vinyals, O.: Neural Discrete Representation Learning, NeurIPS (2018),