[論文レビュー] DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning

Gaoyue Zhou, Haizhou Pan+2 more|arXiv (Cornell University)|Nov 7, 2024

Artificial Intelligence in Games被引用数 5

ひとこと要約

DINO-WM は、事前学習済み DINOv2 パッチ埋め込みを使って潜在空間にオフラインの世界モデルを構築し、専門家デモや報酬設計なしで MPC によるゼロショット計画を可能にする。

ABSTRACT

The ability to predict future outcomes given control actions is fundamental for physical reasoning. However, such predictive models, often called world models, remains challenging to learn and are typically developed for task-specific solutions with online policy learning. To unlock world models' true potential, we argue that they should 1) be trainable on offline, pre-collected trajectories, 2) support test-time behavior optimization, and 3) facilitate task-agnostic reasoning. To this end, we present DINO World Model (DINO-WM), a new method to model visual dynamics without reconstructing the visual world. DINO-WM leverages spatial patch features pre-trained with DINOv2, enabling it to learn from offline behavioral trajectories by predicting future patch features. This allows DINO-WM to achieve observational goals through action sequence optimization, facilitating task-agnostic planning by treating goal features as prediction targets. We demonstrate that DINO-WM achieves zero-shot behavioral solutions at test time on six environments without expert demonstrations, reward modeling, or pre-learned inverse models, outperforming prior state-of-the-art work across diverse task families such as arbitrarily configured mazes, push manipulation with varied object shapes, and multi-particle scenarios.

研究の動機と目的

• オフラインの軌道データからオンライン相互作用やタスク特有の報酬なしで、タスク非依存の世界モデルを学習する。
• 観測をエンコードするために、事前学習済みの視覚パッチ埋め込みを使用し、エンコーダを凍結した状態を保つ。
• 潜在ダイナミクスモデルを訓練して将来のパッチ特徴を予測し、テスト時にはMPCまたは勾配計画を通じて計画を可能にする。
• ピクセルを再構成せず、潜在空間での行動列を最適化して目標観測へ到達する計画を行う。
• さまざまな迷路、操作、およびマルチ粒子シナリオに対する一般化を実証する。

提案手法

• 観測モデルは凍結された DINOv2 パッチ埋め込みを使用して画像を z_t にエンコードする。
• 遷移モデルは因果アテンションを用いた ViT ベースのアーキテクチャで、過去の潜在状態と行動を取り入れて履歴長さ H で z_{t+1} を予測する。
• 行動（利用可能な場合は位置覚も）を各パッチ表現に連結して予測を条件付けする。
• 訓練は潜在的一貫性損失を用いた教師ありフォーシング: L_pred = || p_theta(enc_theta(o_{t-H:t}), phi(a_{t-H:t})) - enc_theta(o_{t+1}) ||^2.
• 任意のデコーダ q_theta は解釈可能性のため z_t から o_t を再構成し、L_rec = || q_theta(z_t) - o_t ||^2 で別個に訓練される。
• 計画はMPCとCEMを用いて、目標への潜在平均二乗誤差を最小化する: C = || hat{z}_T - z_g ||^2。
• 計画中には地真のピクセル再構成は不要で、計画は潜在空間だけで完全に行われる。

Figure 1: We present DINO-WM , a method for training visual models by using pretrained DINOv2 embeddings of image frames (a). Once trained, given a target observation $o_{T}$ , we can directly optimize agent behavior by planning through DINO-WM using model-predictive control (b). The use of pretrain

実験結果

リサーチクエスチョン

• RQ1 DINO-WM は、事前学習済みの視覚表現を用いてオフラインデータセットから効果的に訓練できるのか？
• RQ2 潜在空間の世界モデルは、デモや報酬なしでテスト時に任意の視覚目標へゼロショット計画を可能にするのか？
• RQ3 事前学習済みパッチ埋め込みを使用することは、異なるタスクファミリーや環境構成での一般化にどのように影響するのか？
• RQ4 パッチベースの表現は、操作タスクに必要な空間構造を全体特徴埋め込みと比較してより良く保持するのか？

主な発見

モデル	PointMaze SR	PushT SR	壁 SR	ロープ CD	粒状 CD
IRIS	0.74	0.32	0.04	1.11	0.37
DreamerV3	1.00	0.04	1.00	2.49	1.05
TD-MPC2	0.00	0.00	0.00	2.52	1.21
Ours	0.98	0.90	0.96	0.41	0.26
(empty)	(empty)	(empty)	(empty)	(empty)	(empty)

• DINO-WM は、最も難しいタスクで前例の最先端と比べて視覚再構成指標（LPIPS）を最大56%向上させる高品質な潜在世界モデルを生み出す。
• DINO-WM を用いた任意の目標への計画は、最も難しいタスクで従来手法と比べ成功率が平均45%向上する。
• DINO-WM は、タスクファミリ内の環境変動（例: 異なる迷路レイアウト、オブジェクト形状）に対して高い成功率で一般化する。
• パッチベースのエンコード（DINO Patch）は、正確な空間理解を要する操作タスクで、グローバルエンコーダ（R3M、ResNet、DINO CLS）を上回る。
• 拡散ベース生成モデルと比較して、DINO-WM は計画のためのより物理的に妥当かつ目標指向の予測を提供する。
• デコーダーを用いた解釈性は、環境間でSSIM/LPIPSの比較で有利な結果を示し、頑健な潜在予測を確認する。

Figure 2: Architecture of DINO-WM . Given observations $o_{t-k:t}$ , we optimize the sequence of actions $a_{t:T-1}$ to minimize the predicted loss to the desired goal $o_{g}$ . All forward computation is done in the latent space $z$ . Here $p_{\theta}$ indicates DINO-WM ’s dynamics model, which is