Parameter-Efficient Distributional RL via Normalizing Flows and a Geometry-Aware Cramér Surrogate

기술 요약: NFDRL

문제 제기
분포 강화 학습 (DistRL) 은 기대 보상뿐만 아니라 전체 보상 분포를 모델링하여 불확실성, 위험, 변동성에 대한 더 풍부한 신호를 제공합니다. 그러나 기존 방법들은 표현력과 매개변수 효율성 사이에서 절충을 강요받습니다.

• 범주형 방법 (예: C51): 고정된 이산 지지점에 의존합니다. 효과적이기는 하지만 간결성이 부족하며, 매개변수 수가 원하는 해상도 (원자의 수) 에 비례하여 선형적으로 증가합니다. 또한 보상에 경직된 상한과 하한을 부과하여 동적이거나 두꺼운 꼬리를 가진 보상이 있는 도메인에서의 적용성을 제한합니다.
• 분위수 방법 (예: QR-DQN, IQN): 무제한 보상을 처리하기 위해 분위수 위치를 학습합니다. 그러나 분포를 디랙 델타 (계단 함수) 의 혼합물로 근사합니다. 이 표현은 매끄럽고 다중 모달이거나 두꺼운 꼬리를 가진 분포에 비효율적일 수 있으며, 종종 단순한 밀도를 근사하기 위해 많은 분위수가 필요하고 세밀한 분포 형태를 포착하지 못합니다.
• 연속 밀도 방법: 연속 확률 밀도 함수 (PDF) 를 모델링하면 적응형 지지점과 매끄러움을 제공하지만, 이전 접근법들은 이론적으로 타당성 (벨만 연산자의 수축 보장) 과 실제 최적화 가능성 (편향 없는 샘플 기울기 제공) 을 동시에 만족하는 학습 목적 함수를 정의하는 데 어려움을 겪었습니다.

방법론: NFDRL
저자들은 **연속 정규화 흐름 (continuous normalizing flows)**을 사용하여 보상 분포를 모델링하는 매개변수 효율적인 DistRL 아키텍처인 NFDRL을 제안합니다.

1. 아키텍처:

   • 모델은 기본 노이즈 샘플 $z \sim U$를 보상 결과 $y = f_\theta(z)$로 매핑하는 조건부 정규화 흐름을 사용합니다.
   • 흐름은 **가우시안 혼합 누적 분포 함수 (CDF)**를 통해 구성됩니다. 신경망 $h_\theta$는 각 행동에 대해 상태 $x$를 혼합 CDF 의 매개변수로 매핑합니다.
   • CDF 의 유계성 ($[0, 1]$) 을 해결하기 위해 추가적인 아핀 변환이 출력을 동적 범위 $[-G_{max}, G_{max}]$로 재조정하여, 모델이 고정된 경계 없이 실제 세계 보상의 지지점에 적응할 수 있도록 합니다.
   • 목표 분포는 예측된 흐름을 부트스트랩 함수 ($r + \gamma y$) 와 결합하여 구성하며, 변수 변환 공식을 통해 정규화를 유지하기 위해 가역 흐름 계층으로 처리됩니다.

2. 학습 목적 함수 (기하학적 인식 크라메르 대리 함수):

   • 표준 KL 발산은 겹치지 않는 지지점에는 적합하지 않으며, 정확한 와세르슈타인 최소화는 편향된 기울기로 고통받습니다.
   • 저자들은 원시 밀도가 아닌 확률 질량 위에 정의된 크라메르에서 영감을 받은 기하학적 인식 거리를 도입합니다.
   • 손실 함수는 다음과 같이 정의됩니다:
     $$D(\eta, \hat{T}\eta) = \left( \frac{1}{N^2} \sum_{i,j} (w_i - v_i)^2 \cdot |y_i - y_j| \right)^{1/2}$$
     여기서 $w_i$와 $v_i$는 지지점 $y_i$에서의 확률 질량입니다.
   • 부트스트랩 단계로 인해 예측 지지점과 목표 지지점 사이의 불일치를 처리하기 위해, 이 방법은 손실을 계산하기 전에 공유 지지점에서 밀도를 평가하기 위해 **커널 밀도 추정 (KDE)**을 사용합니다.

주요 기여

1. 매개변수 효율성: 해상도가 증가함에 따라 모델 크기가 커지는 범주형 베이스라인과 달리, NFDRL은 분포의 유효 해상도와 관계없이 일정한 매개변수 footprint을 유지하여 무한한 유효 해상도와 동적 지지점을 제공합니다.
2. 이론적 타당성: 저자들은 질량 기반 거리가 **진정한 거리 (true metric)**이며 관련 분포 벨만 연산자가 $\sqrt{\gamma}$-수축임을 증명합니다. 이는 고유한 고정점으로의 수렴을 보장합니다.
3. 편향 없는 기울기: 제안된 목적 함수는 편향 없는 샘플 기울기를 허용하며, 이는 이전 PDF 기반 DistRL 방법들 (와세르슈타인과 달리 종종 편향된 근사를 필요로 함) 에서는 일반적으로 결여된 특성입니다.
4. 표현력: 흐름의 연속성으로 인해 모델은 분위수 기반 방법들이 종종 흐리게 하거나 해결하지 못하는 복잡한 다중 모달 보상 지형을 포착할 수 있습니다.

결과

• 토이 MDP: 통제된 환경에서 NFDRL 은 매끄러운 PDF 를 가진 풍부하고 다중 모달인 보상 지형 (예: 이모달 분포) 을 성공적으로 복원하는 반면, 분위수 베이스라인 (IQN) 은 모달리티를 흐리게 하는 노이즈가 많고 흐릿한 분포를 생성했습니다.
• Atari-5 벤치마크: NFDRL 은 범주형 베이스라인 (C51) 과 경쟁력 있는 성능을 달성하고 최상위 분위수 방법 (IQN) 의 성능에 근접합니다.
  • Atari-5 벤치마크에서 NFDRL (대리 변형) 은 평균 인간 정규화 점수 418을 기록하여 C51(330) 과 DQN(189) 을 능가했으나 IQN(525) 에는 약간 뒤처졌습니다.
  • 특히 NFDRL-S 는 Double Dunk와 QBert*와 같은 특정 게임에서 IQN 을 능가했습니다.
• 효율성: NFDRL 은 훨씬 적은 매개변수로 C51 의 성능과 일치합니다. C51 의 매개변수 수가 원자 수에 비례하여 선형적으로 증가하는 반면, NFDRL 은 고정된 크기를 유지하여 C51 의 용량을 11 개의 원자만으로 달성합니다.

의의 및 주장
이 논문은 NFDRL 을 Atari 에서 IQN 과 같은 강력한 분위수 기반 에이전트를 대체하는 보편적인 솔루션이 아니라, 명확한 절충안을 가진 간결하고 연속적인 대안으로 위치시킵니다. 그 주요 의의는 일정한 매개변수 footprint 으로 연속 보상 분포에 대한 이론적으로 엄격한 프레임워크 (수축 및 편향 없는 기울기) 를 제공한다는 점에 있습니다. 이는 기대값뿐만 아니라 보상 분포의 상세한 형태 (예: 다중 모달리티, 꼬리 위험) 에 따라 하류 결정이 의존하는 시나리오에서 특히 유용합니다. 저자들은 더 높은 학습 분산, 일부 베이스라인에 비해 느린 수렴, 그리고 KDE 정렬 단계의 계산 비용과 같은 한계를 인정하며, 이를 향후 개선 영역으로 보고 있습니다.