TRAM: Test-Time Risk Adaptation with Mixture of Agents

技術的概要：TRAM（エージェントの混合によるテスト時リスク適応）

1. 問題定式化

本論文は、強化学習（RL）の展開における重要なギャップに対処する：安全仕様は訓練の「後」に変わることが多い。エージェントはソースタスクのセットで訓練されるが、新しい危険マップ、修正されたリスク閾値、または異なる行動整合制約（例：新しい参照行動）を備えた環境に展開される可能性がある。

既存のリスク感受性 RL 手法は、通常、訓練段階中にリスクモデル（分散、CVaR、または特定の制約など）を固定する。このアプローチは以下の理由で脆い：

1. 過度の専門化：特定のリスク代理変数（戻り値の分散など）を用いて訓練することは、ソース方策の多様性を縮小させ、異なる将来のリスク仕様のもとで安全である可能性のある行動を排除する。
2. 不十分なリスク表現：分散ベースのリスクは、空間的危険（バリアへの近接など）、参照からの行動のドリフト、または局所的な変動を捉えられない。特に、軌道の戻り値分散がゼロである決定論的設定において顕著である。
3. 再訓練コスト：新しいリスク仕様への適応は通常、再訓練または微調整を必要とし、これは計算コストが高く、モデルが展開された後はしばしば非現実的である。

本論文は、この問題をゼロ更新の展開時適応として定式化する。目標は、固定されたリスク中立のソース方策のライブラリ $\{\pi_j\}$ を再利用し、ソースパラメータを変更することなく、テスト時に新たに指定されたリスク汎関数 $\rho(\pi)$ を満たすことである。目的は、リスク制約 $\rho(\pi) \le \delta$ の下で期待戻り値 $J(\pi)$ を最大化することであり、ペナルティ付き代理変数を通じて近似される：
$$\max_{\pi} \tilde{J}_c(\pi) = J(\pi) - c \rho(\pi)$$

2. 手法：TRAM

提案手法であるTRAMは、ソース訓練とリスク仕様の分離を行う。これは 2 つのフェーズで動作する：

フェーズ 1：リスク中立ソース訓練
ソース方策は、標準的なリスク中立の目的関数を用いて関連するソースタスク上で訓練される。これにより、多様な「行動の凸包」が維持され、後で再結合可能な広範な占有分布（状態 - 行動訪問パターン）がライブラリに含まれることが保証される。

フェーズ 2：テスト時ソーススコアリングと構成
展開時、TRAM は再訓練を行わない。代わりに、ソーススコア付き構成則に基づいてソースライブラリから行動を選択することで、展開方策 $\pi_T$ を構築する。

1. リスク評価：各ソース方策 $\pi_j$ について、TRAM はキャッシュされた、グローバルな占有ベースのリスクスコア $\rho(d_{\pi_j})$ を計算する。これにより、割引状態 - 行動占有 $d_\pi$ 上で定義された柔軟なリスク汎関数が可能になる。例：
   • バリア曝露：危険領域の近くで費やされた時間。
   • 行動の発散：安全な参照行動からの KL 発散。
   • 局所的変動：占有におけるステップごとの変動。
2. 価値推定：各ソースのターゲットタスク価値 $Q^{\pi_j}_T(s, a)$ を現在の状態 $s$ で推定する。これは、Successor Features（SFs）を用いたドット積 $Q^{\pi_j}_T(s, a) = \psi_{\pi_j}(s, a)^\top w_T$、またはロールアウト/クリティックを通じて効率的に行うことができる。
3. 行動選択：展開される行動は、リスク調整済みソーススコアを最大化することで選択される：
   $$S_j(s, a) = Q^{\pi_j}_T(s, a) - c \rho(d_{\pi_j})$$
   $$\pi_T(a|s) \in \arg\max_{b \in A} \max_{j \in [n]} S_j(s, b)$$
   アクティブなソースは軌道に沿って変化し、現在の状態と特定のリスク仕様に応答する「局所的に継ぎ接ぎされた」方策が作成される。

理論的保証：
本論文は、TRAM を完全な占有制御問題の厳密な最適化ではなく、計算可能なソーススコア付き代理変数として明示的に位置づける。

• 局所化された境界：理論的分析（定理 4.1）は、TRAM 方策と最適リスク中立方策との間の最適化不足の境界を提供する。この境界は、ターゲットと最良の一致ソース間の局所化された報酬の不一致に依存し、さらにリスクペナルティ重み $c$ とリスク汎関数のリプシッツ定数に比例する項を含む。
• ソース凸包の不一致：継ぎ接ぎされた方策の実現された占有は、ソース占有の凸包と異なる可能性があることが論文で認められている。付録 C は、キャッシュされたソースリスクと展開された方策の実現リスクの間のギャップを埋める測定可能な「継ぎ接ぎ不一致」項（$\epsilon_{stitch}$）を導入している。

3. 主要な貢献

1. ゼロ更新適応の定式化：パラメータ更新なしに、固定されたソース方策を新しい占有ベースのリスクに適合させる設定を定式化した。
2. TRAM アルゴリズム：テスト時に多様なリスクタイプ（バリア、発散、変動）をサポートするソーススコア付き構成則を導入し、分散などの固定代理変数に依存する手法と対比させた。
3. 理論的基盤：
   • リスク中立ソース訓練が最悪ケースの展開リスク後悔を最小化し（定理 4.3）、最も広範な再利用可能な占有凸包を維持することを証明。
   • 報酬転移の質と展開リスクペナルティを分離する局所化された転移保証。
   • キャッシュされたソースリスクと実現された方策リスクの間の「継ぎ接ぎ不一致」の分析。
4. 実証的検証：4 つの異なるドメインでの実証：
   • グリッドワールド：分散ベースの手法が失敗する空間的危険やステップごとの変動など、表現力豊かなリスクモデリングの検証。
   • MuJoCo Reacher：目標性能を維持しつつ、バリア違反を削減。
   • Safety-Gymnasium：凍結された安全代理変数や分散ベースのベースラインと比較して、低いバリアヒット確率を達成。
   • LLM 整合：微調整なしに、固定された 7B 言語モデルを新しい行動リスク制約（安全な参照への KL）に適合させ、純粋な報酬最大化よりも改善された整合性を示す。

4. 結果

実験により、TRAM は再訓練を必要とすることなく、展開安全性（危険曝露またはバリアヒットの削減）を継続的に向上させながら、報酬性能を維持することが示された。

• グリッドワールドでは、TRAM は空間的危険を回避し、変化するリスク条件に適応することに成功したが、リスク認識ベースライン（RaSF など）は均一に保守的であったり、空間的リスクを検出できなかったりする。
• Reacherでは、訓練後に新しいバリアが導入された際、TRAM は失敗率を 85% から 26% に削減し、目標までの距離への影響は最小限であった。
• Safety-Gymnasiumでは、TRAM は最も低いバリアヒット確率（ベースラインの 0.502 に対して 0.209）を達成し、競争力のある正規化 CVaR を示した。
• LLM 整合では、TRAM はリスクフリーの転移方策よりも GPT-4 によるペアワイズ比較で好まれる応答を生成し、高次元のテキスト空間における報酬モデルの過剰最適化を防ぐことができることを示した。

5. 意義と主張

本論文は、TRAM が RL システムのためのモジュール型展開レイヤを提供すると主張する。その主な意義は、リスクの決定を訓練段階から展開段階へシフトさせる点にある。

• モジュール性：オペレーターは、ソース行動がすでに学習された後に、安全要件（危険マップ、行動参照など）を指定できる。
• 効率性：テスト時に再訓練や新しい制約付き MDP の解決にかかる計算コストを回避し、代わりに高速なソーススコアリングに依存する。
• 範囲に関する正直さ：著者は主張において謙虚であり、TRAM が代理手法であることを明示的に述べている。これは継ぎ接ぎされた方策の完全な占有制御問題を厳密に解決すると主張するものではない。代わりに、ソース凸包代理変数と実現リスクの間の不一致に対する測定可能な境界を提供し、リスク仕様が動的で再訓練が禁止されているシナリオにおける実用的なヒューリスティックとして機能する。

この研究は、多くの展開シナリオ（ロボティクス、ナビゲーション、LLM など）において、多様でリスク中立な行動ライブラリを維持し、テスト時スコアリングを通じてそれらを構成することが、あらゆる潜在的な将来の制約に対して専門的なリスク認識方策を訓練するよりも堅牢な戦略であることを示唆している。