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) 에이전트는 배포 시 분포 변화 (distribution shifts) 로 인해 실제 환경에서 종종 실패합니다. 에이전트가 훈련 데이터에서 잘 수행되더라도 역학이나 상태 전이의 미세한 차이가 충돌이나 비용 임계값 초과와 같은 안전하지 않은 행동으로 이어질 수 있습니다.

기존 솔루션은 다음과 같은 중대한 한계를 안고 있습니다:

• 신념 기반 적응 (Belief-based adaptation): 불확실성을 모델링하기 위해 처음부터 재훈련이 필요합니다.
• 제약/보수적 RL (Constrained/Conservative RL): 종종 성능 (수익) 을 희생하거나 추가 비용 신호를 통한 미세 조정이 필요합니다.
• 테스트 시간 적응 (Test-time adaptation): 대부분의 현재 방법론은 파라미터 업데이트나 인간 감독 없이 사전 훈련된 오프라인 에이전트를 배포 시 안전 제약에 맞게 적응시킬 수 없습니다.

이 논문이 다루는 핵심 질문은 다음과 같습니다: 오프라인 RL 에이전트가 재훈련 없이 사전 훈련된 분포만을 사용하여 테스트 시간에 안전을 위해 자체 정렬 (self-aligned) 될 수 있는가?

2. 방법론: SAS (안전 위한 자체 정렬, Self-Alignment for Safety)

저자들은 Lyapunov 기반 자체 정렬을 통한 테스트 시간 적응을 가능하게 하는 트랜스포머 기반 프레임워크인 SAS를 제안합니다. 이 방법은 모델 파라미터를 업데이트하지 않고, 대신 **인-컨텍스트 학습 (in-context learning)**과 베이지안 추론에 의존합니다.

핵심 구성 요소

1. 오프라인 RL 을 위한 Lyapunov 밀도 모델 (LDM):

   • 저자들은 오프라인 데이터에서 유도된 **점유 측정 (occupancy measures)**을 사용하여 Lyapunov 안정성을 재구성합니다.
   • 그들은 $E(s, a) = -\log \hat{\rho}_\pi(s, a)$라는 에너지 함수를 정의하는데, 여기서 $\hat{\rho}_\pi$는 상상된 궤적 (imagined trajectories) 에서 추정된 모델 기반 점유 측정입니다.
   • 에이전트가 제어 불변 집합 (control-invariant set)($R_{SAS}^G$) 내에 머무르도록 보장하기 위해 Lyapunov 대리 함수 $G_{SAS}(s, a)$를 구성합니다. 이 집합은 궤적이 저밀도 (안전하지 않거나 분포 밖) 영역으로 진입하지 않도록 보장합니다.
   • $G_{SAS}$는 임의의 궤적을 따라 만나는 최소 에너지로 정의되어, 정책이 안전하지 않고 확률이 낮은 영역으로 붕괴되는 것을 방지합니다.

2. 확률적 추론 관점:

   • 이 논문은 트랜스포머 기반 RL 에이전트를 계층적 RL (Hierarchical RL) 시스템으로 해석합니다. 여기서 고수준 정책은 잠재 기술 ($z_t$) 을 선택하고 저수준 정책은 행동을 실행합니다.
   • 인-컨텍스트 학습은 이러한 잠재 기술에 대한 베이지안 추론으로 프레임됩니다. 프롬프트 (상태 - 행동 쌍의 시퀀스) 를 제공함으로써, 모델은 해당 프롬프트와 일관된 가장 가능성 높은 고수준 정책을 추론합니다.
   • 안전 조건은 관측 가능한 지표 (observable indicators)($U_t$ 및 $V_t$) 로 인코딩됩니다:
     • $U_t$: 현재 상태 - 행동 쌍이 안전 밀도 영역 ($G_{SAS} > 0$) 내에 있는지를 나타냅니다.
     • $V_t$: Lyapunov 함수가 비감소 (non-increasing) 하는지 ($G_{SAS}(s_t, a_t) \geq G_{SAS}(s_{t+1}, a_{t+1})$) 나타냅니다.

3. SAS 알고리즘 (알고리즘 1):
   이 과정은 "상상 (imagination)"과 "선택"의 두 가지 루프로 구성됩니다:

   • 단계 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. 확률적 안전 보장: 저자들은 에이전트가 안전 집합을 벗어날 확률이 상상된 롤아웃 수 ($N$) 와 재상상 시도 수 ($M$) 가 증가함에 따라 지수적으로 감소함을 보여주는 공식적 경계 (정리 4) 를 제공합니다.

4. 실험 결과

이 방법은 Safety Gymnasium, MuJoCo, Bullet-Safety-Gym 벤치마크에서 평가되었습니다.

• 성능 대 안전: SAS 는 Decision Transformer (DT), Constrained Decision Transformer (CDT), 그리고 다양한 안전 RL 베이스라인 (BC-Safe, COptiDICE 등) 과 비교하여 비용과 실패율을 일관되게 감소시켰으며 (최대 2 배 개선), **수익 (보상)**은 유지하거나 향상시켰습니다.
• 일반화: SAS 는 명시적 비용 신호로 훈련되지 않은 백본 (DT) 에 대해서도 안전성을 향상시켰으며, 이는 밀도 기반 Lyapunov 지시가 암시적 안전 대리 변수로 작용함을 보여줍니다.
• 절대화 연구 (Ablation Studies):
  • Lyapunov 하강 ($V_t$): 하강 조건을 제거하면 비용이 크게 증가하여 장기적 안정성을 위해 이것이 필수적임을 입증했습니다.
  • 프롬프트 선택: 무작위 프롬프팅이나 최대 에너지 프롬프팅은 Lyapunov 기반 선택보다 성능이 낮아, 특정 선택 기준의 유효성을 확인했습니다.
  • 모델 강건성: 이 방법은 다양한 프롬프트 길이, 상상된 궤적 수, 그리고 백본 모델 크기 (DT 에서 GPT-2 까지) 에서 강건성을 보였습니다.

5. 중요성과 영향

• LLM 과 RL 의 연결: 이 논문은 대규모 언어 모델 (LLM) 에서 인기 있는 자체 정렬 (self-alignment) 개념을 안전 RL 영역에 성공적으로 적용하여, 에이전트가 자신의 내부 세계 모델을 사용하여 안전에 대해 "추론"할 수 있게 합니다.
• 배포 타당성: 재훈련이나 미세 조정이 필요 없으므로, SAS 는 안전이 중요한 동적이고 불확실한 환경에서 오프라인 RL 에이전트를 확장 가능하게 배포할 수 있는 솔루션을 제공합니다.
• 이론적 엄밀성: Lyapunov 안정성 이론과 베이지안 추론에 기반하여 에이전트가 제어 불변 집합 내에 머무른다는 확률적 보장을 제공함으로써, 휴리스틱 안전 검사를 넘어섰습니다.
• 효율성: 상상 (롤아웃) 을 위한 추론 시간 계산이 필요하지만, 전체 재훈련의 계산 비용을 피하므로 데이터는 정적이지만 환경은 그렇지 않은 시나리오에 적합합니다.

요약하자면, SAS는 Lyapunov 기반 상상을 활용하여 인-컨텍스트 학습을 통해 테스트 시간에 행동을 자체 수정함으로써, 분포 변화에 강건한 오프라인 RL 에이전트를 만드는 중요한 진전을 의미합니다.