Generalization in offline RL: The structure is more important than the amount of pessimism

제목: 오프라인 강화학습에서의 일반화: 비관성의 양보다 구조가 더 중요하다

문제 정의
오프라인 강화학습(Offline Reinforcement Learning, RL)은 추가적인 환경 상호작용 없이 고정된 정적 데이터셋으로부터 최적의 정책을 학습하는 것을 목표로 한다. 분포 외(out-of-distribution, OOD) 행동에 대한 과대평가 편향을 완화하기 위해, 대부분의 최첨단 방법론들은 비관적인 가치 학습(pessimistic value learning) 또는 보수성(conservatism)을 채택한다. 그러나 최근 문헌들은 지나치게 보수적인 태도가 새로운 시나리오로의 일반화나 궤적 스티칭(trajectory stitching)을 저해할 수 있음을 시사한다. 일부 접근 방식은 일반화를 개선하기 위해 "완만한(mild)" 보수성을 옹호하지만, 제로샷 정책 전이(zero-shot policy transfer, ZSPT) 성능을 결정하는 핵심 요인이 비관성의 '정도'인지, 아니면 비관적 가치 함수의 '구조적 특성'인지는 여전히 불분명하다. 특히, 환경에 기저 대칭성(예: 회전 불변성)이 존재하는 경우, 매우 비관적이지만 대칭적인 가치 함수가 완만하지만 비대칭적인 가치 함수보다 더 잘 일반화될 수 있는지에 대해서는 알려진 바가 없다.

방법론
저자들은 컨텍스추얼 MDP(Contextual MDPs) 하의 '일반화를 통한 불변성 ZSPT(Generalization-through-Invariance ZSPT, GTI-ZSPT)' 프레임워크 내에서 이론적 분석과 경험적 검증을 결합하여 이 문제를 조사한다.

1. 이론적 프레임워크:

   • 저자들은 에이전트가 훈련 데이터에서 관찰되는 부분군 $B \leq G$에도 불구하고, 최적해의 올바른 대칭 그룹 $G$를 학습하는 능력을 통해 일반화를 공식화한다.
   • 저자들은 오프라인 RL 메커니즘의 복잡성을 분리하기 위해 **Q-값 증류(Q-value distillation)**를 대리 지표로 활용한다. 이를 통해 특정 비관적 가치 함수을 학습하는 효과만을 격리하여 분석한다.
   • 분석은 무한히 넓은 신경망(infinitely wide neural network) 한계를 가정하며, NTK(Neural Tangent Kernel)와 딥 앙상블의 등변성(equivariance)에 관한 Gerken & Kessel (2024)의 결과를 활용한다.
   • 두 가지 핵심 정리를 증명한다:
     • 정리 1: 만약 비관적인 가치 함수가 훈련 데이터의 기저 대칭 그룹 $B$를 준수한다면, 비관성의 수준($\eta_{max}$)이 얼마나 크더라도 에이전트는 최적의 일반화 성능을 달성할 수 있다.
     • 정리 2: 반대로, 미세하게 비관적이지만 대칭 구조를 위반하는 가치 함수($\hat{Q}_{asym}$)는, 매우 비관적이지만 대칭적인 가치 함수($\hat{Q}_{sym}$)에 비해 일반화 성능이 임의로 훨씬 나빠질 수 있는 사례가 존재한다.

2. 경험적 검증:

   • 환경: 회전 대칭성을 가진 "Rotational Reacher" 연속 제어 환경을 사용한다. 훈련 컨텍스트는 $C_4$ 부분군(90도 회전)을 다루며, 테스트는 전체 $SO(2)$ 그룹을 다룬다.
   • 알고리즘: 본 연구는 Implicit Q-Learning (IQL)과 Conservative Q-Learning (CQL)을 평가한다.
   • 데이터 증강(Data Augmentation, DA) 전략: 표준 DA(증강된 데이터셋에서 훈련)와 데이터 증강 일관성(Data Augmentation Consistency, DAC) 규제(정규화)를 비교한다. DAC는 원래 입력과 증강된 입력 사이의 네트워크 출력(또는 잠재 표현) 차이를 최소화하는 손실 항을 추가함으로써 대칭성을 강제한다.
   • 가설: 저자들은 DAC가 오프라인 RL에 더 우월하다고 주장하는데, 이는 DAC가 가치 추정의 정확도보다 대칭 구조(symmetry)를 명시적으로 강제함으로써 가치 함수의 대칭성을 우선시하기 때문이다. 이는 그들의 이론적 결과와 일치한다.

주요 기여

• 이론적 통찰: 본 논문은 일반화에 있어 비관성의 '양'보다 비관성의 '구조'가 더 중요하다는 것을 입증한다. 비관적 가치 함수가 환경의 대칭성을 준수한다면, 임의로 높은 수준의 비관성 하에서도 최적의 일반화가 가능하다는 것을 증명한다.
• 역설적 발견: 완만한 비관성이 반드시 더 나은 일반화를 보장하는 것은 아님을 보여준다. 만약 완만한 비관성이 요구되는 대칭 구조를 깨뜨린다면, 매우 비관적이지만 대칭적인 대안보다 성능이 임의로 저하될 수 있다.
• 방법론적 제안: 저자들은 데이터를 단순히 증강된 데이터셋으로 훈련하는 표준적인 관행 대신, **정책 추출 과정에서의 일관성 손실(consistency loss during policy extraction)**을 통한 데이터 증강 적용을 제안한다(특히 IQL의 액터에 대해). 이 방식은 정책/가치 함수의 대칭 구조를 직접적으로 강제한다.
• 경험적 증거: IQL 및 CQL에 대한 실험 결과, DAC-Output(네트워크 출력에 일관성을 강제하는 방식)이 가장 우수한 일반화 성능을 보였으며, 이는 증강된 데이터셋 기반 훈련(Aug-D)이나 증강이 없는 베이스라인보다 뛰어난 성능을 나타냈다.

결과

• Q-값 증류: Rotational Reacher를 이용한 이론적 검증에서, 대칭적인 비관적 타겟($\hat{Q}^*_{sym}$)은 비관성이 한 자릿수 증가하더라도 최적에 가까운 테스트 성능을 유지했다. 반면, 비대칭적 타겟($\hat{Q}^*_{asym}$)은 비관성이 증가함에 따라 테스트 성능이 크게 저하되었다.
• 오프라인 RL 실험:
  • IQL과 CQL 모두에서, DAC-Output을 적용했을 때 전문가(expert), 혼합(mixed), 하위 최적(suboptimal) 데이터셋 모두에서 일관되게 가장 높은 테스트 리턴을 기록했다.
  • 표준 데이터 증강(Aug-D)은 증강이 없는 베이스라인보다는 개선되었으나, 일반적으로 일관성 손실 방식보다 성능이 낮거나 비슷했다.
  • IQL에서 (액터 없이) 크리틱(critic)에만 DA를 적용한 경우 베이스라인보다 개선되지 않았으며, 이는 정책 추출 구성 요소에 대칭성을 강제하는 것이 매우 중요하다는 점을 확인시켜 주었다.

의의 및 주장
본 논문은 "보수성을 줄이는 것이 일반화에 필요하다"는 기존의 직관이 불완전하다고 주장한다. 대신, 저자들은 올바른 구조적 대칭성을 강제하는 것이 성공적인 오프라인 RL 제로샷 전이의 핵심 동력이라고 주장한다.

이 연구의 의의는 "얼마나 보수적이어야 하는가?"라는 질문에서 "우리의 보수성이 환경의 대칭성을 존중하는가?"로 초점을 전환했다는 데 있다. 저자들은 데이터 분포가 비대칭적 편향을 유도할 수 있는 오프라인 RL 상황에서 **데이터 증강 일관성(DAC)**과 같은 기술이 필수적이라고 본다. 정책 추출 과정에서 가치 정확도보다 대칭성을 우선시함으로써, DAC는 가치 함수가 매우 비관적인 상황에서도 에이전트가 최적으로 일반화할 수 있도록 한다. 이는 데이터 증강을 단순히 훈련을 위한 데이터 크기를 늘리는 용도로 사용하는 표준 관행에 도전하며, 대신 이를 구조적 불변성을 강제하기 위한 규제 도구로 사용해야 함을 시사한다.

한계점
저자들은 다음과 같은 한계점을 인정한다:

• 이론적 증명은 신경망의 **무한 폭 한계(infinite width limit)**에 의존하고 있으나, 유한 폭 신경망에 대한 경험적 결과가 이를 뒷받침한다.
• 분석은 데이터 분포와 일치하는 **내재적 그룹 대칭성(intrinsic group symmetries)**에 초점을 맞춘다. 규제 목적으로만 사용되는 외재적 또는 불일치하는 변환(예: 무작위 노이즈)에 대해서도 이 결론이 유효한지는 아직 입증되지 않았다.
• 현재 결과는 Rotational Reacher 태스크와 두 가지 특정 알고리즘(IQL, CQL)을 통해 검증되었으며, 더 복잡한 환경과 알고리즘에 대한 광범위한 검증이 필요하다.