Generalization in offline RL: The structure is more important than the amount of pessimism

タイトル：オフライン強化学習における汎化：構造は悲観性の量よりも重要である

問題提起
オフライン強化学習（Offline RL）は、さらなる環境との相互作用を行うことなく、固定された静的なデータセットから最適な方策を学習することを目的としている。分布外（OOD）のアクションに対する過大評価バイアスを軽減するために、最先端の手法の多くは悲観的な価値学習や保守性を採用している。しかし、近年の文献は、過度に保守的になることは、新しいシナリオへの汎化を妨げたり、軌跡のステッチング（stitching）を阻害したりする可能性があることを示唆している。一部のアプローチは、汎化を向上させるために「穏やかな」保守性を提唱しているが、ゼロショット方策転送（ZSPT）の性能を決定付ける主要な要因が、悲観性の「度合い」なのか、それとも悲観的価値関数の「構造的特性」なのかについては依然として不明である。具体的には、基礎となる対称性（例：回転不変性）を持つ環境において、高度に悲観的だが対称な価値関数は、穏やかに悲観的だが非対称な価値関数よりも優れた汎化を実現できるのかどうかは分かっていない。

手法
著者らは、Contextual MDPs（CMDPs）における「不変性を介した汎化（GTI-ZSPT）」フレームワークの下で、理論的解析と経験的検証を組み合わせてこの問いを調査している。

1. 理論的枠組み:

   • 著者らは、訓練データにおいて観察される部分群 $B \leq G$ にかかわらず、最適解の正しい対称群 $G$ を学習できる能力として、汎化を定式化している。
   • 著者らは、オフラインRLメカニズムの複雑さを排除するために、**Q値蒸留（Q-value distillation）**をプロキシとして利用している。これにより、特定の悲観的価値関数を学習することの効果を分離している。
   • 解析では、無限に広いニューラルネットワークの極限を仮定し、Neural Tangent Kernel（NTK）および深層アンサンブルの等変性に関するGerken & Kessel (2024) の結果を利用している。
   • 2つの主要な定理を証明している：
     • 定理1: もし悲観的価値関数が訓練データの持つ対称群 $B$ を尊重しているならば、エージェントは、悲観性のレベル（$\eta_{max}$）がいかに大きくても、テスト環境において最適な汎化性能を達成できる。
     • 定理2: 逆に、穏やかに悲観的ではあるものの、対称構造を逸脱した価値関数（$\hat{Q}_{asym}$）が存在する場合、それは対称な価値関数（$\hat{Q}_{sym}$）と比較して、汎化性能が極端に悪化することが保証されるケースがある。

2. 経験的検証:

   • 環境: 回転対称な「Rotational Reacher」連続制御環境。訓練コンテキストは $C_4$ 部分群（90度回転）をカバーし、テストは完全な $SO(2)$ 群をカバーする。
   • アルゴリズム: 著者らは、Implicit Q-Learning (IQL) と Conservative Q-Learning (CQL) を評価している。
   • データ拡張（DA）戦略: 本論文では、標準的なデータ拡張（拡張されたデータセット上で訓練する手法）と、**データ拡張一貫性（DAC）**正則化を比較している。DACは、元の入力と拡張された入力の間でネットワークの出力（または潜在表現）の差を最小化する損失項を追加することで、対称性を強制する。
   • 仮説: 著者らは、DACがオフラインRLにおいて優れていると主張している。なぜなら、DACは価値関数や方策の構造（対称性）を明示的に強制し、価値の正確性よりも対称性を優先するためであり、これは彼らの理論的知見と一致している。

主要な貢献

• 理論的洞察: 本論文は、汎化においては悲観性の量よりも、悲観性の構造がより重要であることを確立している。悲観的な価値関数が環境の対称性を尊重している限り、どれほど高い悲観性であっても最適な汎化が可能であることを証明している。
• 直感に反する発見: 穏やかな悲観性が必ずしも優れた汎化を保証するわけではないことを示している。もし穏やかな悲観性が要求される対称構造を破壊する場合、高度に悲観的だが対称な代替案と比較して、性能が極端に悪化する可能性がある。
• 手法的な提案: 著者らは、データ拡張を単に拡張されたデータセットでの訓練として用いるのではなく、方策抽出時における一貫性損失（consistency loss）を介したデータ拡張を適用することを提案している（具体的にはIQLのアクターに対して）。このアプローチは、方策や価値関数の対称的な構造を直接的に強制する。
• 経験的証拠: IQLおよびCQLに関する実験において、DAC-Output（ネットワーク出力に対して一貫性を強制する手法）が最高のテストリターンを一貫して達成し、ノーオーグメンテーション（拡張なし）のベースラインや拡張訓練（Aug-D）を上回った。

結果

• Q値蒸留: Rotational Reacherを用いた理論的検証において、対称な悲観的ターゲット（$\hat{Q}^*_{sym}$）は、悲観性が1桁増加しても、ほぼ最適なテスト性能を維持した。対照的に、非対称なターゲット（$\hat{Q}^*_{asym}$）は、悲観性が増すにつれてテスト性能が著しく低下した。
• オフラインRL実験:
  • IQLおよびCQLの両方において、DAC-Outputを適用することで、エキスパート、混合、およびサブオプティマルなデータセットのすべてにおいて、一貫して高いテストリターンが得られた。
  • 標準的なデータ拡張（Aug-D）は、ノーオーグメンテーションのベースラインよりは改善が見られたものの、一般に一貫性損失を用いた手法を下回るか、同等の結果であった。
  • IQLにおいて、DAをクリティック（Critic）のみに適用した場合（アクターには適用しない場合）、ベースラインからの改善は見られず、方策抽出コンポーネントに対して対称性を強制することが極めて重要であることが確認された。

意義と主張
本論文は、「保守性を下げる必要がある」という従来の直感は不完全であると主張している。代わりに、著者らは、正しい構造的対称性を強制することこそが、オフラインRLにおける成功したゼロショット転送の主要な推進力であると論じている。

本研究の意義は、オフラインRLにおける「どの程度保守的であるべきか？」という問いから、「我々の保守性は環境の対称性を尊重しているか？」という問いへと焦点をシフトさせた点にある。著者らは、データセットの被覆率が非対称なバイアスを誘発しやすいオフラインRLにおいて、**データ拡張一貫性（DAC）**のような手法が不可欠であると断言している。方策抽出の際に価値の正確性よりも対称性を優先することで、DACは、基礎となる価値関数が高度に悲観的である場合でも、エージェントが最適に汎化することを可能にする。これは、データ拡張を単に訓練のためのデータサイズを増やすために用いるという標準的な慣行に異を唱えるものであり、むしろ構造的な不変性を強制するための正則化として用いるべきであることを示唆している。

限界
著者らは以下の限界を認めている：

• 理論的な証明は、ニューラルネットワークの無限幅の極限に依存しているが、有限幅のネットワークにおける経験的な結果がこの理論を支持している。
• 解析は、データの分布と一致する固有の群対称性に焦点を当てている。正則化としてのみ用いられる外的な、あるいは一貫性のない変換（例：ランダムノイズ）についても同様の結論が成り立つかどうかは未証明である。
• 知見は現在、Rotational Reacherタスクと2つの特定のアルゴリズム（IQL, CQL）において検証されている。より複雑な環境やアルゴリズムに対する広範な検証が必要である。