Lyapunov-Guided Self-Alignment: Test-Time Adaptation for Offline Safe Reinforcement Learning

以下は、論文「Lyapunov-Guided Self-Alignment: Test-Time Adaptation for Offline Safe Reinforcement Learning」の詳細な技術的概要です。

1. 問題定義

静的なデータセット（例：D4RL、DSRL）で訓練されたオフライン強化学習（RL）エージェントは、分布のシフトにより実世界環境に展開された際に失敗することが多い。エージェントが訓練データ上で良好な性能を発揮していても、動的特性や状態遷移のわずかな違いが、衝突やコスト閾値の超過などの安全でない行動につながる可能性がある。

既存の解決策は重大な限界に直面している：

• 信念に基づく適応: 不確実性をモデル化するために、ゼロから再訓練が必要である。
• 制約付き/保守的な RL: 多くの場合、性能（リターン）を犠牲にするか、追加のコスト信号を用いた微調整を必要とする。
• テスト時適応: 現在のほとんどの手法は、パラメータの更新や人間の監督なしに、事前学習されたオフラインエージェントを展開時に安全制約に適応させることができない。

ここで扱われる核心的な問いは以下の通りである：オフライン RL エージェントは、再訓練なしに、事前学習された分布のみを用いてテスト時に安全性のために自己整合（self-align）できるか？

2. 手法：SAS（安全性のための自己整合）

著者は、Lyapunov 誘導型自己整合によるテスト時適応を可能にするトランスフォーマーベースのフレームワークSASを提案する。この手法はモデルパラメータの更新を行わず、代わりにコンテキスト内学習とベイズ推論に依存して動作する。

中核コンポーネント

1. オフライン RL 用の Lyapunov 密度モデル（LDM）:

   • 著者は、オフラインデータから導出された**占有測度（occupancy measures）**を用いて Lyapunov 安定性を再定式化する。
   • 彼らは、想像された軌道から推定されたモデルベースの占有測度 $\hat{\rho}_\pi$ を用いて、エネルギー関数 $E(s, a) = -\log \hat{\rho}_\pi(s, a)$ を定義する。
   • エージェントが制御不変集合（control-invariant set）（$R_{SAS}^G$）内に留まることを保証するよう、Lyapunov 代理関数 $G_{SAS}(s, a)$ が構築される。この集合は、軌道が低密度（安全でない/分布外）領域に進入しないことを保証する。
   • $G_{SAS}$ は、任意の軌道上で遭遇する最小エネルギーとして定義され、これにより方策が安全でない低確率領域へ崩壊するのを防ぐ。

2. 確率的推論の視点:

   • 本論文は、トランスフォーマーベースの RL エージェントを階層的 RLシステムとして解釈する。ここで、高次方策は潜在スキル（$z_t$）を選択し、低次方策が行動を実行する。
   • コンテキスト内学習は、これらの潜在スキルに対するベイズ推論として枠組み化される。プロンプト（状態 - 行動対のシーケンス）を提供することで、モデルはそのプロンプトと整合する最も可能性の高い高次方策を推論する。
   • 安全条件は観測可能な指標（$U_t$ および $V_t$）として符号化される：
     • $U_t$: 現在の状態 - 行動対が安全な密度領域内（$G_{SAS} > 0$）にあるかどうかを示す。
     • $V_t$: Lyapunov 関数が非増加である（$G_{SAS}(s_t, a_t) \geq G_{SAS}(s_{t+1}, a_{t+1})$）かどうかを示す。

3. SAS アルゴリズム（アルゴリズム 1）:
   このプロセスは、「想像」と選択の 2 つのループを含む：

   • ステップ 1（条件 $U_t$）: 固定された初期状態から、トランスフォーマーは $N$ 個の軌道を想像する。それは、最大エネルギー（最悪ケースのリスク）を最小化する軌道セグメントを選択し、初期プロンプトを形成する。
   • ステップ 2（条件 $V_t$）: 初期プロンプトを用いて、トランスフォーマーは $M$ 個の新しい軌道を想像する。それは、累積 Lyapunov 降下（$V_t$ 指標の和）を最大化する軌道を選択する。
   • ステップ 3（展開）: 選択された安全なセグメントが、トランスフォーマーのためのコンテキスト内プロンプト（$p_{1:K}$）として使用される。エージェントはその後、このプロンプトと現在の状態に基づいて行動を条件付け、勾配更新なしに事前学習された方策を安全な行動へと「誘導」する。

3. 主要な貢献

1. 新しい定式化: 占有測度に基づくオフライン RL 用の Lyapunov 安定性の新たな定義。これにより、明示的なコスト関数ではなく、データ密度を通じて安全制約を強制可能となる。
2. 理論的枠組み: ベイズ推論によるトランスフォーマーベース RL の階層的 RL 解釈。ここでプロンプティングは、潜在的高次スキルの推論に対応する。これは、コンテキスト内学習が安全性において機能する理論的基盤を提供する。
3. テスト時適応手法（SAS）: エージェントの行動を再整合させるLyapunov 誘導型プロンプトを生成する実用的なアルゴリズム。オフラインデータから導出された制御不変集合にエージェントを制限することで、安全性を保証する。
4. 確率的安全性保証: 著者は、定理 4 において、エージェントが安全集合から逸脱する確率が、想像されたロールアウト数（$N$）と再想像試行数（$M$）とともに指数関数的に減少することを示す形式的な境界を提供する。

4. 実験結果

この手法は、Safety Gymnasium、MuJoCo、およびBullet-Safety-Gymのベンチマークで評価された。

• 性能対安全性: SAS は、Decision Transformer (DT)、Constrained Decision Transformer (CDT)、および BC-Safe や COptiDICE などの各種安全 RL ベースラインと比較して、コストと失敗率を一貫して削減（最大 2 倍の改善）し、リターン（報酬）を維持または向上させた。
• 汎化性: SAS は、明示的なコスト信号で訓練されていないバックボーン（DT など）に対しても安全性を向上させた。これは、密度ベースの Lyapunov 誘導が暗黙的な安全性の代理として機能することを示している。
• アブレーション研究:
  • Lyapunov 降下（$V_t$）: 降下条件を除去するとコストが大幅に増加し、長期的な安定性に対するその必要性が証明された。
  • プロンプト選択: ランダムなプロンプティングや最大エネルギープロンプティングは、Lyapunov 誘導型選択よりも性能が劣り、特定の選択基準の有効性が確認された。
  • モデルの頑健性: この手法は、異なるプロンプト長さ、想像された軌道数、およびバックボーンモデルサイズ（DT から GPT-2 まで）において頑健性を示した。

5. 意義と影響

• LLM と RL の架け橋: 本論文は、大規模言語モデル（LLM）で人気のある自己整合の概念を安全 RL の領域に成功裏に適応させ、エージェントが自身の内部世界モデルを用いて安全性について「推論」することを可能にした。
• 展開の実用性: 再訓練や微調整の必要性を排除することで、SAS は安全性が極めて重要である動的で不確実な環境におけるオフライン RL エージェントの展開のためのスケーラブルな解決策を提供する。
• 理論的厳密性: Lyapunov 安定性理論とベイズ推論に基づき、エージェントが制御不変集合内に留まる確率的保証を提供することで、ヒューリスティックな安全性チェックを超えている。
• 効率性: 想像（ロールアウト）のための推論時計算を必要とするものの、完全な再訓練の計算コストを回避するため、データは静的だが環境はそうでないシナリオに適している。

要約すると、SASは、Lyapunov 誘導型想像を活用してコンテキスト内学習を通じてテスト時に行動を自己修正することにより、オフライン RL エージェントを分布シフトに対して堅牢にするための重要な前進を表している。