JEPA世界モデルの計画を改善する非等方的な表現

はじめに

モデルベース強化学習における世界モデルは、次の状態を予測するだけでは十分ではない。予測結果を使って、どの行動列を選ぶべきか判断できなければならない。JEPA型の世界モデルでは、観測を潜在表現に変換し、行動条件付きの予測器で状態遷移を学習する。その後、予測した終端表現と目標表現のユークリッド距離を使って候補行動を順位付けする構成がよく用いられる。

ここで重要になるのが、潜在空間の距離が本当にタスクの良し悪しを表しているのかという点だ。論文「Anisotropic Representations Improve Planning in JEPA World Models」は、この対応関係が自動的には保証されないと指摘する。

正確な予測だけでは足りない理由

表現が正確な予測を支え、表現崩壊も起きていなくても、計画に適した幾何が得られるとは限らない。一般的な等方的ガウス正則化は、潜在空間の各方向を似た尺度で扱う。一方、制御タスクでは、ある方向の誤差が結果に大きく影響し、別の方向の差はそれほど重要でない場合がある。

すべての方向をほぼ同じように重み付けすると、プランナーは実行可能な結果をタスク上のコストとは異なる順序で評価する可能性がある。つまり問題は、モデルが正しく予測できるかだけでなく、予測誤差を意思決定の場面でどのように測るかにもある。

AnisoWMの仕組み

AnisoWMとΛRegは、プランナーを複雑化するのではなく、学習時の目標分布を変更する。固定された等方的ガウス目標の代わりに、学習可能な対角共分散を用い、潜在表現の各方向に異なる分散を割り当てる。さらに、固定トレース制約と非等方性に関する制約を設け、特定の次元だけが極端に拡大・縮小する解を抑える。

予測目的、予測器の構造、そして計画時のユークリッド距離は変更されない。学習目標の統計的な形を変えることで、既存の距離計算がより意味のあるものになるよう表現自体を整える発想である。論文では、予測の難しさや学習データの分布が分散配分に与える影響、さらにその計量が計画後悔を減らす条件も分析している。

意義

4つの視覚制御環境において、AnisoWMはすべてでLeWorldModelより高い計画成功率を示した。また、潜在的な計画コストと実際のタスク結果の一致度も改善した。提供された素材には環境名や詳細な数値は含まれていないため、ここから読み取れる中心的な結論は、改善が新しいプランナーではなく表現幾何の調整から得られたという点である。

この研究は、世界モデルを予測誤差や表現崩壊だけで評価することへの注意を促す。潜在距離が実際の成果を正しく順位付けできるかを検証することが、今後のJEPA型計画システムにおける重要な診断項目になりそうだ。

出典：Hugging Face Daily Papers