JEPA-TTT、世界モデルをテスト時に継続適応させる

背景

世界モデルは、環境の未来状態を予測し、その予測を使って行動計画を立てる。しかし、訓練時と実運用時で摩擦や質量、アクチュエータの応答などが変わると、学習済みの遷移モデルは急速に不正確になる可能性がある。JEPA-TTTは、この問題に対して、事前学習済みの潜在世界モデルをテスト中も更新し続ける仕組みを提案する。

仕組み

JEPA-TTTの土台は、アクション条件付きのJEPA（Joint-Embedding Predictive Architecture）世界モデルだ。モデル全体を再学習するのではなく、環境変化の影響を受けやすい部分に更新を絞っている。

• 潜在ダイナミクス予測器だけを更新：視覚エンコーダーを固定し、テストデータによって事前学習済みの視覚表現が崩れることを抑える。報酬ヘッドも固定し、元のタスク目的を維持する。
• 自己教師ありで適応：エージェントが観測した後続状態を利用して予測を修正するため、オンライン報酬や目標画像を必要としない。
• 密なリプレイを利用：軌跡から一部だけを取り出すのではなく、各時間オフセットで予測ウィンドウを作成し、成長するバッファーに保存する。更新時にはそこからミニバッチをサンプリングする。
• エピソード間で経験を蓄積：各エピソード終了時に適応をリセットせず、テスト中に得た情報を継続利用する。

この設計では、視覚表現やタスクに関する既存知識を保ちながら、潜在空間における未来予測だけを新しい力学に合わせて調整できる。画像を完全に再構成する必要がない点も特徴だ。

結果と意味

評価は4つの連続制御環境と8種類のダイナミクス変化を対象に行われた。JEPA-TTTは、すべての変化条件で固定されたJEPA世界モデルを上回った。テストエピソード500回後には、自己回帰型の潜在予測誤差が平均83%減少し、計画性能は固定モデルに対して153%向上した。

この結果は、世界モデルを「一度学習したら固定して使う」ものから、運用中に校正し続けるものへ拡張できる可能性を示す。オンライン報酬がなくても、観測された状態と予測のずれを適応信号として利用できるからだ。一方で、古い経験と新しい経験をどう配分するか、環境が継続的に変化する場合に忘却をどう防ぐかなど、課題も残る。今回の評価は連続制御環境が中心であり、複雑な視覚タスクや開放的な環境への適用性は今後の検証が必要だ。

出典：Hugging Face Daily Papers