UNIQ: Conformal Calibration for Adaptive Conservatism in Offline Reinforcement Learning

技術要約：UNIQ – オフライン強化学習における適応的保守性のための共形キャリブレーション

1. 問題提起

オフライン強化学習（RL）は、**分布シフト（distribution shift）**という極めて重要な課題に直面している。学習された方策が、固定されたオフラインデータセット内で稀であるか、あるいは存在しない状態行動ペアをクエリしてしまう可能性があるためである。標準的な時間差（TD）学習法は、これらの分布外（OOD）領域において過大な推論を行い、壊滅的な過大評価と方策の崩壊を招く。

既存の解決策は、行動複製（behavioral cloning）による制約、保守的な価値学習（CQL、IQLなど）、およびアンサンブルベースの不確実性の3つのカテゴリに概ね分類される。Implicit Q-Learning (IQL) は、固定されたエキスタイルのパラメータ $\tau$ を用いてOODの過大評価を抑制することで、計算負荷が低く安定したベースラインを提供しているが、構造的な限界を抱えている。それは、単一の固定された $\tau$ をすべての状態に対して一様に適用してしまう点である。

• データが密集している領域では、この固定された保守性は不必要に制限的であり、潜在的な性能向上の機会を逃している。
• データが疎な領域では、固定された $\tau$ では過大評価を防ぐのに不十分な場合がある。

本論文は、EDACのような大規模なアンサンブル手法のような膨大な計算コストをかけることなく、データの被覆率が低い場所では保守性を強め、被覆率が高い場所では緩和するという、**状態ごとの（per-state）**適応的な保守性メカニズムが必要であると主張している。

2. 手法：UNIQ

UNIQ (Uncertainty-Informed Quantile) は、固定されたエキスタイル $\tau$ を、共形予測（conformal prediction）によって校正された不確実性に駆動される状態適応型の $\tau(s)$ に置き換えることで、IQLのバックボーンを拡張したものである。本手法は、以下の3つのコアコンポーネントで構成される。

A. マルチ・エキスタイル・バリュー・アンサンブル
単一の価値ネットワークを訓練する代わりに、UNIQは3つの固定されたエキスタイルレベル $\bar{\tau} \in \{0.5, 0.7, 0.9\}$ に適合させた $N_v$ 個の価値ネットワークのアンサンブルを訓練する。これにより、$3N_v$ 個の価値ヘッドが得られる。このマルチレゾリューションな適合により、2つの相補的な不確実性信号を抽出する：

1. 認識的不確実性 ($\sigma_{ens}$): アンサンブルの予測における $\tau=0.7$ レベルでの標準偏差であり、クリティック間の不一致を捉える。
2. アレオリックな広がり ($\Delta_\tau$): $\tau=0.9$ における平均予測と $\tau=0.5$ における平均予測の差であり、リターンの固有の確率性を捉える。

B. 分割共形キャリブレーション（Split Conformal Calibration）
生のアンサンブルの不一致 $\sigma_{ens}(s)$ は、タスクに依存しスケールに敏感である。これを分布に依存しない不確実性スコアへと正規化するために、UNIQは分割共形予測を採用する：

• 学習データから校正用スプリット $D_{cal}$ を保持する。
• 各校正遷移に対して非適合度スコア $\alpha_i$ を計算し、アンサンブルの $\tau=0.7$ の価値関数のベルマン残差を測定する。
• これらのスコアの $(1-\delta)$ 分位点 $\hat{q}$ を計算し、有限サンプルでの被覆保証を持つデータ駆動型の閾値を提供する。
• 正規化された不確実性 $u(s)$ は次のように計算される：
  $$u(s) = \frac{\sigma_{ens}(s)}{\hat{q} + \epsilon}$$
  ここで、$u(s) > 1$ はその状態の被覆が不十分であることを示し（不一致が校正された閾値を超える）、$u(s) < 1$ は被覆が良いことを示す。

C. 状態適応型保守性
正規化された不確実性 $u(s)$ は、シグモイド・スケジュールを介して適応的なエキスタイル $\tau(s)$ にマッピングされる：
$$\tau(s) = \tau_{min} + (\tau_{max} - \tau_{min}) \cdot \sigma_{sig}(-\beta_\tau(u(s) - 1))$$

• 高不確実性 ($u(s) \gg 1$): $\tau(s) \to \tau_{min}$ となり、より保守的（低い）な価値推定を強制する。
• 低不確実性 ($u(s) \ll 1$): $\tau(s) \to \tau_{max}$ となり、より楽観的な価値推定を許容する。

さらに、UNIQはQ関数の更新のために、グローバルな悲観的価値ターゲット $V_{pcss}(s) = \bar{V}(0.7)(s) - \kappa \sigma_{ens}(s)$ を使用する。ここで $\kappa$ はタスクごとに調整される。

3. 主な貢献

• メカニズムの革新: UNIQは、共形キャリブレーションを用いて不確実性推定をスケーリングすることにより、IQLのエキスタイル・パラメータを分布に依存しない方法で状態依存的にする初のモデルフリーな手法を導入した。
• 効率性: EDACやSAC-Nのような、クリティック数に比例して計算量が増大し大量のVRAMを必要とするアンサンブル重視の手法とは異なり、UNIQはIQLに近いメモリコストで動作する。これは、共有された低ランクの価値ヘッドと軽量な共形キャリブレーション・ステップを使用することで実現されている。
• 透明性: 著者らは、性能が劣るケースを明示的に報告しており、本研究を、あらゆる手法に対する絶対的なSOTA（最先端）の支配を主張するものではなく、パフォーマンスと効率性の境界における実用的なメカニズムの貢献として位置づけている。

4. 実験結果

実験は D4RL MuJoCo ベンチマーク（HalfCheetah、Hopper、Walker2dの9つのタスク）を用いて行われた。

• IQLとの比較: UNIQは全9タスクにおいてIQLを上回り、IQLの平均スコア81.6に対し、正規化平均スコア85.2を達成した。
  • 向上幅: 改善は特に Hopper と Walker2d タスクにおいて顕著であり、特に medium-replay データセットにおいて顕著である（例：Walker2d-medium-replayで+7.2、Hopper-medium-replayで+4.2）。
  • HalfCheetah: 向上はわずか（+0.1から+1.5）であり、滑らかでよく被覆されたダイナミクスは適応的な保守性の恩恵をあまり受けにくいという仮説と一致している。
• 重いアンサンブルとの比較:
  • UNIQ（平均85.2）は、集計スコアにおいてEDAC（平均92.9）およびReBRAC（平均89.7）には及ばない。
  • しかし、UNIQは大幅に低い計算コストでこれを達成している。
• 効率性:
  • VRAM: UNIQのピークは ~250 MB であり、IQLの530 MB、EDACの2500 MB、ReBRACの~1200 MBと比較して非常に低い。
  • 計算量: UNIQは、EDAC (N=50) よりもステップあたり約4.6倍高速であり、IQLよりはわずか1.1倍遅いだけである。

5. 意義と主張

本論文は、UNIQを、重いアンサンブル手法を凌駕する普遍的なSOTAソリューションとしてではなく、IQL型の手法における固定されたグローバルな保守性という構造的なボトルネックに対処するための実用的なメカニズムとして位置づけている。

• 標的を絞った改善: 本手法は、不均質な被覆率を持つ環境（例：挙動モードが急激に変化するリプレイが豊富なデータセット）において特に効果的である。このようなシナリオにおいて、被覆が豊かな領域では保守性を緩和し、疎な領域では維持するというUNIQの能力は、大きな利益をもたらす。
• リソースへのアクセシビリティ: 近いIQLのメモリコストで動作することにより、UNIQは、膨大なオーバーヘッドを負うことなく、限られた計算リソース（例：シングルGPUやMIGインスタンス）を持つ実務者がIQLよりも高い性能を得るための実行可能な道筋を提供する。
• 今後の方向性: 著者らは、状態依存的な悲観係数 $\kappa(s)$ の自動選択、および共形キャリブレーションの早期ウォームスタートの必要性を、主要な次のステップとして挙げている。

要約すると、UNIQは、分布に依存しない不確実性キャリブレーションによって導かれる適応的な保守性が、性能と効率性の間で価値の高いトレードオフを提供し、大規模な計算スケーリングを必要とせずに、複雑で不均質なデータ領域におけるオフライン強化学習をより堅牢にすることを示している。