Distributional Inverse Reinforcement Learning

1. 문제 정의 (Problem Definition)

기존의 역강화학습 (IRL) 은 주로 관찰된 전문가의 행동 (Trajectories) 을 통해 결정론적인 (Deterministic) 보상 함수를 복원하거나, 기대 보상 (Expected Return) 만을 일치시키는 데 초점을 맞추고 있습니다. 그러나 실제 세계의 많은 환경 (로봇 조작, 신경과학적 행동 등) 에서 보상 신호는 본질적으로 확률적 (Stochastic) 입니다.

• 기존 방법의 한계:

  • 대부분의 IRL 방법 (MaxEntIRL 등) 은 보상을 단일 값 (Point Estimate) 으로 가정하여, 보상의 분산 (Variance) 이나 고차 모멘트 (Higher-order moments) 와 같은 불확실성 정보를 무시합니다.
  • 베이지안 IRL (BIRL) 은 보상 파라미터에 대한 사후 분포를 추정하지만, 여전히 기대 보상을 최적화하는 프레임워크에 의존하므로 보상의 전체 분포 구조를 포착하지 못합니다.
  • 결과적으로 학습된 정책은 보상의 평균에만 민감하고, 위험 (Risk) 이나 변동성에 둔감하여 실제 위험 민감형 (Risk-sensitive) 환경에서 실패할 수 있습니다.

• 핵심 문제:

  • 관찰된 행동 데이터로부터 보상 함수의 전체 확률 분포 (Full Reward Distribution) 를 직접 학습하고, 이를 기반으로 분포를 인지하는 정책 (Distribution-aware Policy) 을 복원하는 것은 어떻게 가능할까?

2. 방법론 (Methodology)

저자들은 DistIRL (Distributional Inverse Reinforcement Learning) 이라는 새로운 프레임워크를 제안합니다. 이 프레임워크는 보상 함수와 회수 (Return) 의 분포를 동시에 모델링합니다.

2.1. 1 차 확률적 우세 (First-Order Stochastic Dominance, FSD) 기반 학습

기존의 기대 보상 일치 대신, FSD (First-Order Stochastic Dominance) 개념을 도입하여 전체 분포를 일치시키려 합니다.

• 목표: 전문가의 회수 분포 ($Z_E$) 가 학습된 에이전트의 회수 분포 ($Z_\pi$) 를 1 차 확률적으로 우세하게 (FSD) 하도록 유도합니다 ($Z_E \succeq_{FSD} Z_\pi$).
• 손실 함수: FSD 위반 정도를 측정하는 적분 형태의 손실 함수를 정의합니다.
  $$L_{FSD} = \int_{-\infty}^{\infty} [F_{Z_E}(z) - F_{Z_\pi}(z)]_+ dz$$
  여기서 $F$는 누적 분포 함수 (CDF) 이며, $[\cdot]_+$는 양의 부분만 취합니다. 이는 전문가의 분포가 학습된 분포보다 항상 우위에 있어야 함을 의미합니다.
• 에너지 기반 모델 (EBM): 이 FSD 위반 손실을 에너지 함수로 간주하여, 보상 분포 $q_\phi(r|s,a)$ 에 대한 가능도 (Likelihood) 를 정의하고 변분 추론 (Variational Inference) 을 통해 사후 분포를 근사합니다.

2.2. 분포 인식 정책 학습 (Distribution-aware Policy Learning)

보상 분포가 학습되면, 이를 바탕으로 위험 민감형 정책을 학습합니다.

• 변형 위험 측정 (Distortion Risk Measure, DRM): FSD 조건을 직접 최적화하는 것은 계산적으로 어렵기 때문에 (지시 함수가 미분 불가능), 변형 위험 측정 (DRM) 을 사용하여 근사합니다.
• 목적 함수:
  $$\max_\phi \int_0^1 F^{-1}_{Z_\pi}(v) d\tilde{\xi}(v) + H(\pi_\phi)$$
  여기서 $F^{-1}$은 분위수 함수 (Quantile function) 이고, $\tilde{\xi}$는 변형 함수 (Distortion function) 입니다. 이 함수는 분포의 특정 부분 (예: 하위 분위수) 에 더 큰 가중치를 주어 위험 회피 (Risk-averse) 또는 위험 추구 (Risk-seeking) 행동을 유도할 수 있습니다.
• 알고리즘:
  1. 보상 네트워크: FSD 손실을 최소화하도록 보상 분포 파라미터를 업데이트.
  2. 크리틱 (Critic): 분위수 회귀 (Quantile Regression) 를 사용하여 분포적 가치 함수를 학습.
  3. 정책 (Policy): DRM 기반 목적 함수를 최대화하도록 정책을 업데이트.

3. 주요 기여 (Key Contributions)

1. 보상 분포 학습 (Reward Distribution Learning):
   • 오프라인 IRL 설정에서 보상을 결정론적 함수가 아닌 확률 변수로 모델링하고, FSD 목적 함수를 통해 평균 이상의 고차 모멘트까지 포착하는 최초의 체계적인 프레임워크를 제안했습니다.
2. 분포 인지 정책 학습 (Distribution-aware Policy Learning):
   • 학습된 보상 분포를 활용하여 위험 민감형 정책을 복원하며, 이는 모방 학습 (Imitation Learning) 을 위험이 있는 시나리오에 적용할 수 있는 기반을 제공합니다.
3. 이론적 분석 (Theoretical Analysis):
   • 제안된 알고리즘의 수렴성을 분석하여, $\mathcal{O}(\epsilon^{-2})$ 반복 복잡도 (Iteration Complexity) 를 가짐을 증명했습니다.
4. 실증적 검증 (Empirical Validation):
   • 합성 데이터, 실제 신경 행동 데이터 (마우스 도파민 신호), MuJoCo 제어 태스크에서 SOTA(State-of-the-Art) 성능을 입증했습니다.

4. 실험 결과 (Results)

• Gridworld (격자 세계):
  • 불확실한 보상을 가진 환경에서 위험 회피형 전문가의 행동을 학습했습니다. 기존 BIRL 은 평균 보상만 맞추는 반면, DistIRL 은 보상의 분산 (Variance) 을 정확히 복원하여 위험 회피 행동을 성공적으로 모방했습니다.
• 마우스 자발적 행동 (Neuroscience Data):
  • 마우스의 자발적 행동 데이터와 도파민 신호를 분석했습니다. 도파민 신호는 본질적으로 비대칭적이고 확률적이므로, 기존 결정론적 모델은 한계가 있었습니다.
  • 결과: DistIRL (특히 왜도 정규 분포를 사용한 S-DistIRL) 은 도파민 신호의 전체 분포 형태 (왜도, 첨도 등) 를 가장 정확하게 복원했으며, 도파민 변동성과 학습된 보상 간의 상관관계가 가장 높았습니다. 이는 행동 데이터만으로 내부 신경 보상 구조를 추론할 수 있음을 시사합니다.
• MuJoCo (D4RL 벤치마크):
  • Risk-sensitive D4RL: 안전 관련 조건 (충돌, 속도 제한 등) 이 확률적 페널티로 주어지는 환경에서 평가했습니다.
  • 결과: 기존 오프라인 IRL 방법 (ValueDICE, Offline ML-IRL 등) 과 비교하여 DistIRL 이 훨씬 높은 평균 회수 (Return) 를 달성했습니다. 특히, BIRL 은 평균만 맞추려 하여 위험 민감한 행동 (안전한 경로 선택) 을 학습하지 못했으나, DistIRL 은 이를 성공적으로 학습했습니다.
  • Risk-neutral D4RL: 결정론적 보상 환경에서도 경쟁력 있는 성능을 보여 방법론의 일반화 능력을 입증했습니다.

5. 의의 및 중요성 (Significance)

• 실제 세계 적용성 향상: 로봇 조작, 자율 주행, 생물학적 행동 모델링 등 보상이 불확실하고 위험이 존재하는 실제 환경에서 IRL 의 적용 가능성을 크게 확장했습니다.
• 신경과학적 통찰: 행동 데이터로부터 뇌의 보상 신호 (도파민) 의 확률적 특성을 복원할 수 있음을 보여주어, 역강화학습이 신경과학 연구 도구로서도 중요한 가치가 있음을 증명했습니다.
• 위험 민감형 AI: 단순히 "평균적으로 좋은" 행동을 넘어, "위험을 피하는" 또는 "특정 위험 수준을 수용하는" 행동을 학습할 수 있는 강력한 프레임워크를 제공합니다.
• 이론적 엄밀성: FSD 와 DRM 을 결합한 새로운 최적화 프레임워크에 대한 수렴성 분석을 제공하여 방법론의 신뢰성을 높였습니다.

요약하자면, DistIRL은 기존 IRL 이 간과했던 "보상의 불확실성"을 핵심 요소로 삼아, 보상의 전체 분포를 학습하고 이를 통해 위험을 고려한 더 견고하고 인간 (또는 생물) 과 유사한 행동을 모방할 수 있는 새로운 패러다임을 제시합니다.