JEPA と LeWM：学習なしで長期タスクを解決する新手法

世界モデルへの展望：生成対統合埋め込み（JEPA）

現在、AI 分野で「世界モデルとは何か」について議論が進んでいます。多くの注目はピクセル単位の予測を行う「生成的世界モデル」に向けられていますが、ヤン・リークン氏らの研究者らはこれを資源の浪費だと指摘しています。

• 生成モデルの問題点:
  • 行動決定に不可欠でないピクセル情報を多く含む。
  • そのような詳細な予測は計算コストが高く、精度向上に限界がある。
• 提唱される代替案：JEPA (統合埋め込み予測アーキテクチャ)
  • モデルの仕組みは以下の 3 つから構成されます。

コンポーネント	機能	詳細
エンコーダ	観測変換	カメラ画像などの観測を、シーンの特徴を記述する**学習済みエンベッディング（ベクトル）**に変換します。
予測器	次状態予測	エンベッディングとアクションを入力とし、次のフレームのエンベッディングを予測します。
計画器	最適化選択	予測器内で数百通りのアクション順序を試み、目標画像のエンベッディングに最も近い終末状態を選びます。

• 利点: プロセス全体がエンベッディング空間で行われるため、モデルは行動決定に関連する重要な部分だけを学習させればよく、計算効率が高いです。

---

LeWM と Push-T 基準での実績

今年、チームは JEPA 型世界モデルである LeWM を公開し、ロボット制御のベンチマークPush-Tで性能を評価しました。

• Push-T タスク:
  • ロボットが灰色の T 字型ブロックをプッシュするシミュレーション。
  • 目標: ブロックを緑色のターゲット位置に収めること。
  • 環境: シミュレータ上で物理法則が完璧に再現される設定と、学習済みの予測器を使用する設定があります。

実験結果の概要

• LeWM の初期状態では、予測精度に限界があったため長期的なタスク（75 ステップ）での成功率は**22%**程度でした。
• シミュレータ（実物理法則）自身を用いた場合は**42%**と向上しました。
• しかし、予測精度の低い学習済みモデルでも、計画手法を改善することで**85%**という高い成功率を達成しました。

---

4 つの追加手法による性能向上

LeWM を再学習することなく、以下の 4 つのアプローチを取り入れることで、長期的タスクにおける成功確率は **22%から85%**へと大幅に改善されました。

1. サブゴールごとの階層計画 (Hierarchical Planning)

モデルの予測範囲は約 25 ステップまでです。75 ステップを一括で予測しようとすると誤差が累積しますが、区切ります。

• 戦略: 長期的な目標（75 ステップ）を、25 ステップごとのサブゴールに分割します。
• 効果:
  • 一括処理の場合：成功率 22%
  • サブゴール分割の場合：成功率 80% に改善。
• 理由: 短い区間（25 ステップ）はモデルが正確に予測できる範囲内にあるため、誤差が累積しにくくなります。

2. コスト付き潜在メモリーグラフ (Latent Memory Graph)

タスクのデモデータから構築されたグラフ構造を利用します。

• 仕組み:
  • デモ中の各フレームを状態としてリンクし、グラフを形成。
  • 計画器は現在地から目標への経路を検索し、経路上の状態をサブゴールとします。
• 課題と解決:
  • ジャンプ（異なるデモ間の移動）にコストがかからない場合、物理的に不可能な経路が選定されました。
  • 修正: ジャンプに対してコストを課す。
• 結果: 成功率が 62% から 85% に向上し、サブゴールを利用した計画器と同等となりました。

3. トランスフォームされた記憶によるギャップ補完

目標付近の状態データが不足している場合でも、対称性を利用してデータを拡張します。

• Push-T の対称性: テーブル全体のシフトや回転、鏡像処理は物理法則に影響しません。
• 手法: 既存のメモリーグラフからデモ区間を抜き出し、シフト・回転・鏡像して目標付近の状態を作成し、グラフに埋め込みます。
• 効果:
  • メモリ不足による成功率の低下（93%→32%）を防ぎ、**85%**まで回復。
  • 限られたメモリーを有効活用します（1% のデモ＋変換で 10% のグラフと同等）。

4. コスト関数の設計最適化（プッシャーの位置含む）

評価基準に「ロボットの腕（プッシャー）」を含めることで、より現実的な動作計画が可能になります。

• 問題: プッシャーの位置もコスト対象外だと、ブロックを目標地点に持っていっても、ロボット自体が別の場所に移動してタスク失敗と判定されることがありました。
• 解決: コスト計算にブロックとプッシャーの両方を含めるようにします。
  • プッシャー無視時：成功 9/50
  • プッシャー含む場合：成功 39/50 へ大幅改善。
• 拡張性: 予測評価時に、目標画像内のプッシャー位置をバリエーション（16 パターン）で比較し、最も近い一致を選ぶことで、再学習なしで成功率が倍増しました。

---

まとめと今後の課題

現在の達成事項

• 発見された 4 つの手法は、モデルの重みを変更せず、既存の LeWM に直接適用可能です。
• シミュレータ自体に近い予測精度を持つ JEPA モデルでも、計画方法の変更だけで上記の結果を得られます。

今後の課題：デモなしでの汎化能力

• 未解決問題: デモデータが存在しない領域や、デモの流れに逆らう目標（フロー逆向き）への対応は依然として難易度が高いです。
  • デモのないゴールへの成功率はまだ限界があります。
• 方向性: これからの研究では、「デモを超えて独自のサブゴールを提案できる上位層」の構築が鍵となります。
• 現時点では、学習済み階層型モデルである H-JEPA がフラットな LeWM よりも高い性能を示していますが、さらなる改善余地があります。

今後の詳細については、X (旧 Twitter) @haddyalchaer で更新予定です。