기술 요약: 신뢰 영역 Q-어쥔트 매칭 (TRQAM)

문제 정의

본 논문은 사전 훈련된 흐름 정책 (pretrained flow policies) 에 대한 **오프-정책 미세 조정 (off-policy fine-tuning)**의 과제를 다룹니다. 흐름 매칭 정책은 기존의 단일 모드 가우시안 정책보다 풍부한 다중 모드 행동 분포를 제공하지만, 다단계 디노이징 과정을 통한 암시적 정의로 인해 그라디언트 기반 정책 개선이 불안정합니다. 샘플링 체인을 통한 직접 역전파는 계산 비용이 매우 높고 불안정성에 취약합니다.

기존 해결책들은 다단계 동역학을 무시하고 가산 잔차 (additive residuals) 를 학습하거나, 고정된 정책의 표현력에 의해 제한되는 노이즈 공간에서 작동함으로써 이를 우회하려 시도합니다. 최근 **어쥔트 매칭을 통한 Q-학습 (Q-learning with Adjoint Matching, QAM)**이라는 접근법은 미세 조정을 기억이 없는 확률적 최적 제어 (SOC) 문제로 재형성하여, 학습된 크리틱 (critic) 을 어쥔트 매칭을 통해 샘플링 과정을 안내하도록 합니다. 그러나 저자들은 QAM 에 근본적인 취약점이 있음을 규명했습니다: 오프-정책 설정에서 학습된 크리틱은 필연적으로 불완전합니다. 크리틱의 작은 근사 오차는 정책 업데이트가 고정된 "온도" (또는 역온도 $\beta$) 에 의존할 때 기하급수적으로 증폭되어, 사전 훈련된 사전 분포로부터 파괴적인 이탈을 초래하고 결국 모델 붕괴를 야기합니다. 이 불안정성은 그라디언트 클리핑이 있더라도 지속됩니다.

방법론: 신뢰 영역 Q-어쥔트 매칭 (TRQAM)

TRQAM 은 신뢰 영역 메커니즘을 SOC 샘플링 동역학에 직접 통합하여 안정적인 오프-정책 미세 조정 알고리즘을 제안합니다. 핵심 혁신은 미세 조정된 정책과 사전 훈련된 사전 분포 간의 경로 공간 쿨백 - 라이블러 (KL) 발산을 적응적으로 제어하는 데 있습니다.

주요 이론적 통찰

1. 오차의 기하급수적 증폭: 저자들은 고정된 온도를 가진 어쥔트 매칭이 크리틱 오차의 기하급수적 증폭을 겪는다는 것을 (Lemma 1) 공식화했습니다. 고정된 $\beta$는 신뢰할 수 있는 크리틱을 활용함과 동시에 신뢰할 수 없는 크리틱으로부터 보호하는 것을 동시에 달성할 수 없습니다.
2. $\lambda$를 신뢰 영역 파라미터로: TRQAM 은 SOC SDE 의 확산 계수를 $\sqrt{\lambda}$만큼 스케일링하는 스칼라 파라미터 $\lambda$를 도입합니다.
3. 정확한 KL 항등식 (Theorem 1): 기르사노프 (Girsanov) 정리를 사용하여, 저자들은 확산을 $\sqrt{\lambda}$만큼 스케일링하면 경로 공간 KL 발산이 $\lambda$의 정확한 폐쇄형 함수가 됨을 증명합니다:
   $$D_{KL}(P_u \parallel P_{base}) = \mathbb{E}_{X \sim P_u} \left[ \frac{1}{2\lambda} \int_0^1 \|u(X_\tau, \tau)\|^2 d\tau \right]$$
   이는 파라미터 $\lambda$와 기본 정책으로부터의 편차 사이에 직접적인 구조적 연결을 확립합니다.
4. 종단 KL 상한 (Proposition 1): 경로 공간 KL 은 종단 행동 분포 KL 을 상한으로 묶으므로, 경로 공간 KL 을 제어함으로써 최종 정책 편향을 효과적으로 제한할 수 있음을 보장합니다.

알고리즘 프레임워크

강한 크리틱 그라디언트에 의해 무시될 수 있는 외부 손실 패널티로 KL 항을 추가하는 대신, TRQAM 은 $\lambda$를 샘플링 동역학 내부에 내재화합니다:

• 샘플링: 제어된 SDE 는 확산 계수 $\sqrt{\lambda}\sigma(\tau)$를 사용합니다. $\lambda$를 조정하면 SDE 자체가 재구성되어 제어된 동역학을 기본 동역학에 더 가깝게 또는 더 멀리 끌어당깁니다.
• 적응: 알고리즘은 이산화된 OT 기억 없는 샘플러에서 유도된 몬테카를로 대용치를 통해 실현된 경로 공간 KL($\hat{D}_n$) 을 추정하여 $\lambda$를 동적으로 적응시키기 위해 **투영된 쌍대 하강 (projected dual descent)**을 사용합니다.
• 업데이트 규칙: $\lambda$는 목표 KL 예산 $\epsilon_{KL}$을 강제하도록 업데이트됩니다:
  $$\lambda_{n+1} \leftarrow \max\{0, \lambda_n + \eta_\lambda (D_n - \epsilon_{KL})\}$$
  실현된 KL 이 예산을 초과하면 $\lambda$가 증가하여 동역학을 더 보수적으로 만들고, 예산 미만이면 $\lambda$가 감소하여 더 공격적인 개선을 허용합니다.

주요 기여

1. 취약성 규명: 본 논문은 고정된 온도를 가진 어쥔트 매칭에서 크리틱 오차의 기하급수적 증폭을 규명하고 공식화하여, Robomimic 과 같은 벤치마크에서 모델 붕괴를 일으키는 경향을 입증했습니다.
2. 이론적 연결: 확산 계수를 $\sqrt{\lambda}$만큼 스케일링하면 제어 비용과 경로 공간 KL 사이에 정확한 항등식이 생성되어 $\lambda$를 원칙적인 신뢰 영역 파라미터로 만든다는 것을 증명했습니다.
3. 알고리즘 설계: TRQAM 은 신뢰 영역 파라미터를 SOC 동역학 내부에 내재화하고 쌍대 하강을 통해 이를 적응시키며, KL 상한을 소프트 손실 패널티가 아닌 샘플링 수준에서 강제하는 방법으로 제안되었습니다.
4. 경험적 우위: 광범위한 실험은 TRQAM 이 기존 방법들이 실패하는 곳에서 안정적인 수렴을 달성하고 최첨단 베이스라인을 능가함을 보여줍니다.

실험 결과

저자들은 OGBench벤치마크와 Robomimic조작 벤치마크의 50 개 작업에서 TRQAM 을 평가했습니다.

• 오프라인 RL 성능: OGBench 에서 TRQAM 은 **68%**의 전체 성공률을 달성하여, 가장 강력한 베이스라인 (DSRL, 46%) 과 이전 어쥔트 매칭 방법들 (QAM 35%, QAM-E 45%) 을 크게 능가했습니다. 이러한 이득은 장기 계획 및 조합적 작업에서 특히 두드러졌습니다.
• 안정성: Robomimic 에서 고정된 온도 방법들 (QAM, QAM-E) 은 다양한 하이퍼파라미터 설정에서 $10^{20}$까지 치솟는 어쥔트 손실과 성공률이 거의 0 으로 붕괴되는 현상을 보였습니다. 반면 TRQAM 은 모든 테스트된 예산에서 안정적으로 유지되었습니다.
• 메커니즘 검증: 제거 실험 (Ablation studies) 은 다음을 확인했습니다:
  • 적응이 필수적입니다: 적응형 KL 예산이 고정된 예산보다 성능이 우수합니다.
  • 내재화가 필수적입니다: SDE 내부에 $\lambda$를 내재화하는 것 (TRQAM) 은 KL 예산을 엄격하게 추적하는 반면, 외부 KL 정규화는 실현된 KL 이 목표보다 훨씬 멀리 이탈하게 하여 성능 저하를 초래합니다.
• 민감도: 목표 KL 예산 $\epsilon_{KL}$은 예측 가능한 제어 노브 역할을 합니다; 일반적으로 더 엄격한 예산이 더 나은 성능을 제공하지만, 최적 값은 작업 구조 (예: puzzle-4x4 와 같이 더 큰 상태 공간을 가진 작업은 더 큰 예산이 유익함) 에 따라 체계적으로 변합니다.

중요성과 주장

본 논문은 TRQAM 이 흐름 정책에 대한 기존 오프-정책 미세 조정 방법에 대한 원칙적이고 안정적인 대안을 제공한다고 주장합니다. 기르사노프 정리를 통해 신뢰 영역 제약을 샘플링 동역학에 내재화함으로써 TRQAM 은 QAM 과 같은 방법에서 불완전한 크리틱으로 인해 발생하는 "파괴적인 이탈"을 피합니다.

저자들은 TRQAM 을 온-정책 RL 을 위한 신뢰 영역 정책 최적화 (TRPO) 와 근접 정책 최적화 (PPO) 의 오프-정책 대응물로 위치시킵니다. 그들은 가치 있는 행동적 사전 분포로 작용하는 사전 훈련된 흐름 정책의 경우, 정확한 KL 항등식을 통한 제어된 편향을 유지하는 것이 파기적 망각과 모델 붕괴를 방지하는 데 필수적이라고 주장합니다. 이 방법은 오프라인 RL 과 온라인 미세 조정으로의 전환 모두에 대한 강력한 솔루션으로 제시되며, 안정성을 희생하지 않고 사전 훈련된 기술을 활용하는 신뢰할 수 있는 메커니즘을 제공합니다.