Simulus: Combining Improvements in Sample-Efficient World Model Agents

技術概要：Simulus – サンプル効率型世界モデルエージェントにおける改善点の統合

問題提起
世界モデル（WM）は、サンプル効率型強化学習（RL）の最前線を表していますが、有望な改善点を統合する際の複雑さにより、その潜在能力はしばしば実現されません。個々の進展は存在するものの、それらを組み合わせるには多大な専門知識とエンジニアリング努力が必要です。さらに、既存のトークンベース世界モデル（TBWM）は、画像観測と離散行動に大きく制限されており、連続制御やマルチモーダル設定への適用性が制限されています。加えて、エピステミック不確実性の低減のための内在的動機付けや、優先度付きリプレイなどのコンポーネントは、統合コストの高さと、内在的探索がタスクと無関係な経験に希少な相互作用予算を浪費するリスクがあるため、サンプル効率型エージェントにおいて頻繁に見落とされています。この状況は、個別に検証された改善点が体系的に統合される前の、深層 RL における Rainbow 以前の時代を反映しています。

手法
著者は、一連の相補的な改善点を統合・検証するために設計されたモジュール型トークンベース世界モデルエージェント「Simulus」を導入します。このエージェントは REM アーキテクチャを基盤としており、4 つの明確なモジュール、すなわち表現モジュール（$V$）、世界モデル（$M$）、コントローラー（$C$）、およびリプレイバッファーで構成されています。

1. 柔軟なトークン化フレームワーク（$V$）： Simulus は、観測と行動のモダリティの任意の組み合わせをサポートするモジュール型トークン化システムを採用しています。画像、連続ベクトル、カテゴリカル変数、およびカテゴリカル変数の画像のようなグリッドである「2D カテゴリカル」をネイティブに処理します。エンコーダーは固定長のトークンシーケンスを生成し、これらが統合された表現に連結されます。これにより、エージェントは連続制御や記号的マルチモーダルタスクを含む多様な環境で動作できます。
2. 世界モデル（$M$）： 動力学モデルは、離散トークン上のシーケンスモデリングとして環境の動力学をモデル化するために、並列観測予測（POP）機構を備えた Retentive Network（RetNet）を利用します。
   • 内在的動機付け： エピステミック不確実性を低減するため、Simulus は次の観測予測ヘッドのアンサンブルを採用します。これらのヘッド間の不一致は、Jensen-Shannon 発散（JSD）を介して測定され、内在的報酬（$r_{int}$）を生成します。これらは外発的報酬と組み合わされ、コントローラーがモデル内での「想像」（計画）中に不確実性の高い領域へ向かうよう導きます。
   • 優先度付きリプレイ： リプレイバッファーは、保存された例の世界モデル損失値を維持します。トレーニングは、単一のパラメータ $\alpha$ によって制御される一様サンプリングと優先度付きサンプリングの混合を利用し、サンプルはソフトマックス変換された損失に比例して選択されます。
3. コントローラー（$C$）： コントローラーは、世界モデルによって生成されたシミュレーション経験のみで動作する拡張されたアクター・クリティックアーキテクチャです。
   • 回帰としての分類（RaC）： 無制限の報酬およびリターン信号を頑健に処理するため、Simulus は値予測のために指数関数的ビンを持つ分類ネットワークを採用します。従来の平均二乗誤差ではなく、クロスエントロピーに基づく HL-Gauss 損失手法を利用することで、値推定の安定性を向上させます。
4. トレーニングサイクル： エージェントは、データ収集、表現学習、世界モデル学習、そして想像内での制御学習というサイクルに従います。コントローラーは、$\lambda$-リターンと値ベースラインを備えた REINFORCE 目的関数を用いてトレーニングされます。

主要な貢献

• 見落とされていたコンポーネントの統合： 本論文は、内在的動機付け、優先度付き世界モデルリプレイ、およびハイブリッド RaC アプローチを特定し、統合します。これらのコンポーネントは個々にも効果的であるだけでなく、組み合わせることで相乗的な利益をもたらすことを実証します。
• 多用途なトークンベースエージェント： マルチモーダリティトークン化をサンプル効率型レジームに適応させることで、Simulus は TBWM を、画像のみの離散行動モデルの制限を超え、視覚的、連続的、およびマルチモーダル入力を処理できる多用途エージェントとして確立します。
• モデル誘発型探索の検証： 本研究は、世界モデル誘発型探索が、厳しい相互作用予算の下でもパフォーマンスを大幅に向上させることを示し、サンプル効率型設定において内在的動機付けがコスト高すぎるという考え方に挑戦しています。
• オープンソース基盤： 著者は、コードとモデル重みをオープンソース化することで、TBWM 研究のための強力な出発点を提供し、競争力のある世界モデルエージェントの開発への参入障壁を下げました。

結果
Simulus は、厳格な相互作用予算を設けた 3 つの多様なベンチマークで評価されました。

• Visual Atari 100K： Simulus は、計画不要な手法の中で最先端（SOTA）のサンプル効率を達成し、DreamerV3、TWM、STORM、DIAMOND などのベースラインを上回りました。特に、これは計画不要の世界モデルとして初めて、人間のレベルのインタークォンタイル平均（IQM）および中央値スコアを達成し、26 個のゲームのうち 13 個で人間を超えるパフォーマンスを達成しました。
• 連続制御（DMC Proprioception 500K）： 連続制御タスクにおいて、Simulus は DreamerV3 と一貫して同等かわずかに上回るパフォーマンスを示し、コンパクトな語彙を持つ連続ドメインにおけるトークンベースアーキテクチャの有効性を証明しました。
• 記号的マルチモーダル（Craftax-1M）： 記号的グリッドと連続的特徴を組み合わせた部分的に観測可能な 2D 生存環境において、Simulus は同時代の世界モデル手法およびすべてのモデルフリーベースライン（RND、PPO-RNN、E3B を含む）を上回り、長いトークンシーケンスを効率的に処理しました。

アブレーション研究により、各コンポーネントが個別に寄与することが確認され、特に内在的報酬が極めて重要であることが示されました。すべてのコンポーネントの組み合わせが最も強力なパフォーマンスをもたらしたことは、顕著な相乗効果を示しています。

重要性
本論文は、世界モデルエージェントの潜在能力が実現されていない理由は、モジュール型改善点を組み合わせる難しさにあると主張しています。DQN における Rainbow フレームワークからインスピレーションを得て、Simulus は、トークン化の柔軟性、内在的動機付け、優先度付きリプレイ、および高度な値予測技術の体系的な組み合わせが、最先端のサンプル効率をもたらすことを実証しています。この研究は、コンポーネントが単一構造アーキテクチャでしばしば見られる干渉を避けるよう慎重に統合される限り、トークンベース世界モデルが多様な RL ドメインにとって多用途でスケーラブルなアプローチであることを確立しています。著者は、トークンベースエージェントが現在、いくつかのベースラインと比較してトレーニング速度の制限に直面しているものの、そのモジュール設計はポリシー推論と将来の研究のスケーラビリティにおいて明確な利点を提供すると主張しています。