JEPA世界モデルの因果ダイナミクス情報崩壊を防ぐAVL-JEPA、外乱下の操作成功率を大幅改善

JEPA世界モデルが高次元の視覚情報を保持しながら行動に起因する物理的状態変化の情報を失う「因果ダイナミクス情報崩壊」を防ぎ、外乱下でも制御計画に有用な潜在表現を学習できるか。

JEPAの予測目的は将来潜在の予測可能性に基づいて情報を重み付けするため、視覚的な持続因子が優先され、行動が引き起こす物理的変化の情報が潜在空間で選択的に失われる。その結果、異なる物理的帰結をもたらす2つの行動がほぼ同じ潜在遷移に写像されることがある。既存の大域的な非崩壊正則化は潜在統計の縮退を防ぐが、この選択的な因果ダイナミクス情報の喪失は防げなかった。

従来手法は、VICRegやSIGRegなどの分散・共分散・分布正則化、EMA教師ネットワーク、逆ダイナミクス補助タスクなどでJEPAの表現崩壊を防いできた。しかしこれらは潜在空間の大域的な非縮退や予測可能性を保つことに主眼があり、行動条件付きの物理的変化を区別する情報が選択的に失われる因果ダイナミクス情報崩壊には対処できなかった。

AVLは2段階の訓練を行う。第1段階の行動グラウンディング経路では、実行された行動列を補助ダイナミクスアンカーとして用い、状態のみのゼロ遷移ベースラインを差し引いた遷移帰属行動ヘッドを潜在遷移から学習する。行動ヘッドはLinear-LayerNorm-GELUの3層トランクで、出力はタスク依存である。損失は潜在予測MSE、SIGReg、行動SmoothL1の和である。第2段階の視覚不変性経路では、現在観測のみにガウスノイズを加え、クリーン観測と摂動観測から予測した将来潜在を同一のクリーン目標へ整合させる。このとき行動ヘッドは凍結したまま行動損失を逆伝播させ、情報を捨てるショートカットを防ぐ。

4つのシミュレーション制御タスクで10シード、各50エピソードのクローズドループ計画を評価した。クリーン環境ではAVLの成功率はTwoRoom 93.2%、PushT 96.0%、Reacher 85.2%、Cube 78.2%であり、CubeではLeWMの73.2%から5.0ポイント改善した。ガウスノイズσ=0.10ではPushTの成功率が15.3%から81.5%へ、TwoRoomが65.7%から88.7%へ向上し、彩度シフト+0.35ではPushTが14.0%から62.0%へ、TwoRoomが54.8%から86.7%へ向上した。TwoRoomの物理的帰結整合Spearman相関は0.35から0.53へ、正規化リグレットは0.152から0.061へ改善し、クリーン・ノイズ間のSpearman相関は0.89から0.98へ向上した。

著者らは、AVLがLeWMのSIGReg正則化に依存しており、評価がシミュレーションのみであることを限界として挙げている。また、実ロボット実験は行われておらず、コード公開についても論文では報告されていない。訓練にはNVIDIA RTX 5090を使用し、PushTで1エポック41.7分など一定の計算コストを要する。