技術的サマリー：認識論的モンテカルロ木探索

問題定義

AlphaZero/MuZero（A/MZ）ファミリーのアルゴリズムは、価値と環境ダイナミクスの学習モデルをモンテカルロ木探索（MCTS）と統合することで大きな成功を収めてきた。しかし、重大な限界が存在する。学習モデルは認識論的不確実性（限られたトレーニングデータの網羅性に起因する不確実性）を導入するが、標準的な MCTS は探索プロセス中にこの不確実性の伝播を考慮しない。その結果、A/MZ はスパース報酬環境における深層探索のために MCTS を効果的に活用できない。深層探索とは、エージェントが現在の状態からの距離に関わらず、新規な遷移へと自らを導く能力を必要とするものであり、報酬がスパースで状態空間が広大なアルゴリズム設計やプログラミングなどのタスクにおいて不可欠である。認識論的不確実性を考慮しない場合、探索は不正確なモデル予測に基づく亜最適方策に収束し、状態空間の必要な領域を探索できなくなる可能性がある。

手法：認識論的 MCTS（EMCTS）

著者は、深層探索を促進するために認識論的不確実性を MCTS プロセスに統合する理論的に動機付けられたフレームワークである**認識論的 MCTS（EMCTS）**を提案する。この手法は以下の 3 つの主要な構成要素を含む。

1. 不確実性を伴う探索の定式化

著者は、学習された環境モデル $\hat{M}$ を確率変数としてモデル化する。そして、学習モデル内の価値予測の分散に基づき、最適価値関数 $Q^*$ に対する上界信頼区間（UCB）を導出する。

• 理論的基盤: 定理 1 は、学習モデル $\hat{M}$ に対して、真の最適価値 $Q^*(s, a)$ が、モデル内の最大期待値に、その価値の標準偏差に比例する項（信頼パラメータ $\delta$ によってスケーリングされた）を加えたもので上界付けられることを示している。
• 探索方策: 標準的な PUCT（Predictor Upper Confidence Bound）選択方策は、**認識論的 P/UCT（EP/UCT）**へと修正される。選択基準は以下のようになる。
  $$a = \arg\max_{a} \left( q_{\hat{M}}(s, a) + \beta \sqrt{V[q_{\hat{M}}(s, a)]} + \text{探索項} \right)$$
  ここで、$q_{\hat{M}}$ は推定された価値を、$\sqrt{V[q_{\hat{M}}]}$ は認識論的不確実性を表す。ハイパーパラメータ $\beta$ は、利用と探索のトレードオフを制御する。

2. 認識論的不確実性の伝播

核心的な貢献は、価値だけでなく、探索木全体を通じて不確実性を伝播させるメカニズムにある。

• バックアップにおける不確実性: バックアップステップ $\nu$ の不確実性は、即時報酬の分散と割引された将来価値の不確実性の分散を合計することで計算される。
• ノード価値の不確実性: A/MZ は計画全体を通じて同じモデルを使用するため、バックアップ帰還は相関を持つ。独立性を仮定することを避けるため、著者は個々のバックアップ帰還の標準偏差の和を用いたノード価値 $V[q_{\hat{M}}(s, a)]$ の分散の上界を提案する。
  $$V[q_{\hat{M}}(s, a)] \leq \left( \frac{1}{N(s, a)} \sum_{i=1}^{N(s, a)} \sqrt{V[\nu_i(s, a)]} \right)^2$$
• 推定量: この手法は、報酬（例：ランダムネットワーク蒸留（RND）またはハッシュベースの計数）および価値（例：不確実性ベルマン方程式（UBE））に対する既存の不確実性推定量を利用する。未観測の遷移については、分散を有界確率変数として取り得る最大分散に設定する。

3. 学習された遷移モデルの扱い

理論的導出は既知の遷移モデルを仮定しているが、著者は学習された遷移ダイナミクス（MuZero の場合など）の課題に対処する。彼らは「最大限楽観的」な近似を提案する。これは、軌道内で最初の不確実な遷移に遭遇した時点で、その軌道内のすべての後続予測が最大の不確実性を持つと仮定するものである。これにより、探索目的において UCB が有効な上界であり続けることが保証される。

主要な貢献

1. 認識論的 MCTS（EMCTS）: 学習された価値および/または報酬モデルから認識論的不確実性を推定・伝播するよう MCTS を拡張した新規アルゴリズム。これにより、探索プロセスが不確実な領域を能動的に追求することが可能になる。
2. 理論的フレームワーク: 学習モデルの分散に理論的に根ざした UCB ベースの探索方策（EP/UCT）の導出。これにより、深層探索のための形式的なメカニズムが提供される。
3. 実装: AlphaZero エージェントと組み合わせた並列化された JAX 実装の EMCTS。アセンブリ言語の subleq 環境と Deep Sea ベンチマークに適用された。

実験結果

著者は、EMCTS を 2 つの困難なスパース報酬ドメインで評価した。

1. Subleq プログラミングタスク

• タスク: 特定の関数（正数の否定と恒等関数）を解決するために、subleq アセンブリ言語でコードを記述する。これは約 $16^{10}$ の状態を持つ状態空間の探索を伴う。
• 結果: AlphaZero と組み合わせた EMCTS（E-AZ）は、ベースラインの AlphaZero を大幅に上回った。E-AZ は、ベースラインよりもはるかに少ないサンプル数で、より困難な「恒等関数」タスクを解決した。適切な不確実性推定量（例：IO ハッシュ対完全状態ハッシュ）を使用することが、さらにサンプル効率を向上させることが示された。

2. Deep Sea ベンチマーク

• タスク: エージェントがスパース報酬を伴う一意の最適軌道を見つける必要があるグリッドワールド環境。ランダム探索によって解を見つける確率は、グリッドサイズとともに指数関数的に減少する。
• 結果:
  • 深層探索: ベースラインの A/MZ エージェントは、合理的なトレーニング予算内で Deep Sea の変種（決定論的および確率的報酬の両方）を解決できなかった。対照的に、EMCTS エージェント（E-AZ および E-MZ）はこれらのタスクを解決し、環境サイズに対するサンプル複雑性の亜指数関数的なスケーリングを示した。
  • 探索の利益: 不確実性を行動選択に使用したが、その不確実性を推定するために探索を使用しなかったアブレーション（A/MZ+UBE）に対して、EMCTS は大幅に優位だった。これは、探索そのものが不確実性推定の質を向上させ、より効率的な探索につながることを確認するものである。
  • 頑健性: この手法は、MuZero の学習された遷移ダイナミクス（値等価抽象化）を使用する場合や、確率的報酬が存在する場合でも有効であった。

意義と主張

本論文は、EMCTS がモデルベース強化学習における根本的なギャップ、すなわち標準的な MCTS が探索のために認識論的不確実性を利用できないという問題を解決すると主張している。不確実性の伝播を探索木に統合することで、この手法は A/MZ エージェントに以下を可能にする。

• スパース報酬環境におけるサンプル効率の大幅な向上。
• ベースラインの A/MZ では実質的に解決不可能な難易度の高い探索ベンチマーク（Deep Sea など）の解決。
• 価値予測に対するより良い不確実性推定を提供することによる、オフライン RL やオフポリシー目標生成における信頼性の向上。

著者は、EMCTS を A/MZ ファミリーに対する実用的かつ理論的に動機付けられた強化として位置づけ、深層探索が不可欠であるアルゴリズム設計やスパース報酬を含む実世界応用において、これらのアルゴリズムをより適切に装備させるものであるとしている。