What Must Generalist Agents Remember?

論文タイトル： ジェネラリスト・エージェントは何を記憶すべきか？
著者： Khurram Yamin, Namrata Deka, Maitreyi Swaroop, Albert Ting, Jeff Schneider, Bryan Wilder (カーネギーメロン大学、ジョージア工科大学)

問題提起

本論文は、ジェネラリスト強化学習（RL）エージェントの開発における根本的な問い、すなわち、複数の環境（ドメイン）や多様な目標に対して最適に近い行動をとるために、エージェントは内部情報の何をメモリに保持しなければならないのか、という問題に取り組んでいる。

エージェントに対して現在のタスクやドメインの明示的なインデックス（識別子）やラベルが与えられることを前提とする先行研究とは異なり、本論文では、エージェントがそのようなインデックスを受け取らない設定を研究している。エージェントは自身の軌跡（状態と行動の履歴）のみを観測し、適切な行動選択に必要な潜在的なドメインの同一性を推論しなければならない。著者らはこれを、隠れた変数である固定されたドメイン・インデックス $\sigma$ が遷移ダイナミクスを変化させる一方で、報酬メカニズムは全ドメインで共有されている、潜在ドメイン付きの部分観測マルコフ決定過程（POMDT）として定式化している。

核心となる課題は、異なる2つのドメインが同じ観測（状態）を提示する可能性がある一方で、それらが互いに相容れない最適な行動を要求する場合である。本論文は、メモリが必要となるのはどのような条件下においてか、また、十分に有能なメモリがあれば、エージェントは環境の局所的な遷移ダイナミクスを再構成できるのか、という問いを投げかけている。

手法および理論的枠組み

著者らは、離散時間における目標の集合 $\Psi$ とドメインの集合 $\Sigma$ を用いて相互作用を定式化している。彼らは主に2つの理論的結果を導入している。

1. 表象の分離（定理1）
第1の結果は、メモリの必要条件を確立するものである。著者らは、2つの異なるドメインが、同じ状態 $s^\dagger$ において、大きな価値ギャップによって隔てられた互いに素な近最適行動の集合を要求する「共有ボトルネック（shared bottleneck）」シナリオ（仮定1）を定義している。

• 主張： エージェントがすべての目標とドメインに対して一様に近最適であるならば、その内部メモリ状態 $M_t$ は、このボトルネック状態においてこれらのドメインを区別しなければならない。
• メカニズム： この定理は、ドメインが与えられたときのボトルネックにおけるメモリ状態の条件付き確率分布が、非ゼロの全変動（Total Variation, TV）距離を持つことを証明している。具体的には、エージェントのポリシーが $\epsilon$-近最適である場合、2つのドメイン間のメモリ分布のTV距離は、価値ギャップ、ボトルネックに到達する確率、および割引因子の関数によって下限が抑えられる。
• 示唆： メモリは、標準的なPOMDPにおけるベルーフ状態のような最適制御のための十分統計量であるだけでなく、潜在的なドメインの同一性が行動に関連する場合に、学習された表現が均一な能力を達成するための「必要条件」である。

2. デコーダーの存在（定理2）
第2の結果は、分離に必要なメモリが、環境のダイナミクスを再構成するために十分であるかどうかを調査している。

• 設定： 著者らは、補助的な目標の集合 $\Psi_F$ である「プローブ目標（probe goals）」を導入しており、これらは1ステップの予測テストとして機能する。彼らは「価値十分なメモリ（value-sufficient memory）」を、これらのプローブ目標に対して最適な行動価値を正確に予測できるメモリと定義している。
• 主張： もしメモリが、局所的な遷移カーネルを区別できる（積分確率指標を介して）一連のプローブ目標に対して価値十分であるならば、局所的な介入的遷移カーネル $P^\sigma(\cdot | s, a)$ を近似的に再構成できるデコーダーが存在する。
• 示唆： 豊かな目標のセットに対して近最適なプランニングをサポートするメモリ状態は、本質的に、環境から明示的なワールドモデルが提供されなくても、環境の局所的なダイナミクスをデコードするために十分な情報を含んでいる。

実験による検証

著者らは、ForkWorldと呼ばれるカスタム環境を用いて、これらの理論的主張を検証している。

• 環境： 廊下が分岐点（ボトルネック）へと続くT字路のグリッドワールド。ドメインは潜在的である。 「ノーマル」ドメインでは「上」への移動は北へ、「下」への移動は南へとなるが、「スワップ（入れ替え）」ドメインでは、これらのマッピングが反転している。目標（上または下に到達すること）は観測可能だが、ドメインは観測できない。
• ボトルネック： 分岐点において、最適な行動は隠れたドメインに完全に依存する。メモリを持たないエージェントはドメインを区別できず、推測を強いられるため、成功率は50%に制限される。
• エージェント： 本研究では、メモリレス・エージェント（DQN）、明示的な履歴ウィンドウを持つエージェント（Stacked DQN）、および再帰型エージェント（DRQN）を、ドメインを知っているオラクル（神託）ベースラインと比較している。

主な結果：

1. メモリの必要性： メモリレス・エージェントは50%のベースラインを超えることができず（約19%を達成）、メモリによるドメインの分離なしにはポリシーが安定した解に収束できないことを裏付けた。メモリを備えたエージェント（Stacked DQNおよびDRDRQN）は、75〜88%の成功率を達成し、オラクルとの差を縮めた。
2. 表象の分離： エージェントの内部状態からドメインを予測するように訓練された線形プローブを用いた結果、メモリを備えたエージェントは、分岐点への後続の訪問において、ほぼ完璧なドメイン分類（精度 >97%）を達成した。一方、メモリレス・エージェントはチャンスレベル（50%）に留まった。これは定理1を経験的に裏付けるものである。
3. デコード可能なダイナミクス： 著者らは、エージェントのメモリ表現から次状態を予測するデコーダーを訓練した。メモリを備えたエージェントは、ベースラインと比較して遷移予測誤差を大幅に減少させた。一方、メモリレス・エージェントには改善が見られなかった。これは定理2を支持しており、プランニングに使用されるメモリが局所的な遷移ダイナミクスをエンコードしていることを示している。

主な貢献と意義

本論文は以下の貢献を行っている。

1. メモリの形式的な必要性の提示： 潜在的なドメイン内で動作するジェネラリスト・エージェントにとって、メモリはオプションではなく、ドメインが観測を共有しながらも異なる行動を要求する場合において、一様な近最適性を達成するための必要条件であることを、制御理論的な証明によって示した。
2. ワールドモデルの基盤としてのメモリ： プローブ目標に対する価値予測をサポートするのに十分なほど豊かなメモリ表現は、本質的に局所的な遷移ダイナミクスのデコードをサポートすることを示す。
3. 分析の転換： タスク・インデックスが与えられていることを前提とする従来の「グッド・レギュレーター（良き調節器）」や因果的RLの結果とは異なり、本研究は「履歴からメモリへのマッピング」を分析しており、エージェントが軌跡データのみから潜在的なドメインを暗黙的に学習しなければならないことを示している。

意義：
著者らは、これらの結果が、ジェネラリスト・エージェントにとってのメモリを、ドメインの曖昧さの解消、遷移モデルの再構成、およびプランニングを支える基盤として特徴づけていると主張している。この研究は、潜在的なコンテキストを推論し、局所的なダイナミクスを再構成する能力は、後付けのアーキテクチャ上の特徴ではなく、シフトする一連の環境に対して優れたパフォーマンスを発揮するエージェントにとっての根本的な要件であることを示唆している。成功するジェネラリストRLエージェントは、潜在的なドメインを分離する表現を内部的にエンコードする必要があり、このエンコーディングが環境の因果構造の再構成を自然に促進することを、これらの知見は示唆している。