Efficient Hypergradient Descent for Inverse Reinforcement Learning

기술 요약: 역강화학습을 위한 효율적인 하이퍼그레이디언트 경사하강법 (Efficient Hypergradient Descent for Inverse Reinforcement Learning)

문제 정의
역강화학습(Inverse Reinforcement Learning, IRL)은 전문가의 시연으로부터 보상 함수를 복구하여, 해당 보상 하에서의 최적 정책이 관찰된 행동을 재현하도록 하는 것을 목표로 한다. 저자들은 최대 가능도 IRL(Maximum-Likelihood IRL, ML-IRL)에 집중하며, 이는 이중 레벨 최적화 문제(bilevel optimization problem)로 정식화될 수 있다. 내부 레벨은 학습된 보상 $r_\phi$ 하에서 엔트로피 정규화된 정책을 최적화하고, 외부 레벨은 유도된 정책과 전문가 데이터 사이의 불일치(특히 KL 발산)를 최소화한다.

이 이중 레벨 문제를 암시적 미분(implicit differentiation)을 통해 해결할 때 발생하는 주요 계산 병목 현상은 하이퍼그레이디언트(hypergradient)를 계산하는 과정이다. 이는 내부 목적 함수의 역-헤시안-벡터 곱(inverse-Hessian-vector product)을 요구한다. 헤시안을 명시적으로 구축하는 것은 이차 메모리($O(d_\theta^2)$)를 필요로 하며, 공액 기울기법(conjugate gradients)과 같은 반복적 근사법은 계산 비용이 많이 들고 조건수(conditioning)에 민데히 민감할 수 있다. 또한, 자연 하이퍼그레이디언트 경사하강법(Natural Hypergradient Descent, NHGD)이 헤시안을 피셔 정보 행렬(Fisher information matrix) 대리물로 대체하기 위해 제안되었으나, 표준 NHGD는 고정된 데이터 분포를 가정한다. ML-IRL에서 내부 목적 함수는 정책에 의해 유도된 궤적 분포에 대한 역 KL 발산(reverse KL divergence)을 포함하므로, NHGD를 직접 적용하는 것은 까닭이 쉽지 않다.

방법론
본 논문은 기하학적 구조를 활용하고 매트릭스 스케칭(matrix sketching)을 채택함으로써 ML-IRL을 위한 하이퍼그레이디언트를 효율적으로 계산하는 방법을 제안한다.

1. 피셔-헤시안 등가성 (Fisher-Hessian Equivalence): 저자들은 정확한 내부 최적성과 실현 가능성(realizability, 정책 클래스가 전문가 분포를 완벽하게 모델링할 수 있을 만큼 충분히 풍부한 경우) 하에서, 내부 목적 함수의 헤시안이 할인된 궤적 피셔 정보 행렬($F_\theta$)에 비례함을 증명한다. 즉, $\nabla^2_\theta L_{inner} = \alpha F_\theta$이다. 이를 통해 하이퍼그레이디언트에 필요한 역-헤시안-벡터 곱을 역-피셔-벡터 곱(inverse-Fisher-vector product)으로 대체할 수 있다.
2. 하이퍼그레이디언트 추정: 이 방법은 하이퍼그레이디언트 계산을 세 가지 구성 요소로 분해한다:
   • 정책 파라미터에 대한 외부 기울기(outer gradient).
   • 피셔 행렬과 외부 기울기를 포함하는 댐핑된 선형 시스템(damped linear system)의 해.
   • 보상 기울기와 선형 시스템의 해를 포함하는 혼합 미분 항.
3. 스트리밍 스펙트럴 스케칭 (Streaming Spectral Sketching, SCFD): 밀집된 피셔 행렬을 저장하는 데 드는 $O(d_\theta^2)$ 메모리 비용을 피하기 위해, 저자들은 SCFD 스케치를 사용하여 필요한 역-피셔-벡터 곱을 근사한다. 스코어 행렬을 실제로 구현하는 대신, 알고리즘은 에이전트 궤적으로부터 가중치가 부여된 정책-스코어 벡터를 크기가 $m$인 압축된 스케치로 스트리밍한다. 이는 저장 복잡도를 $O(md_\theta)$로 줄여준다. 이후 스케치는 근사적인 댐핑 선형 시스템을 푸는 데 사용된다.

주요 기여

• 이론적 유도: 저자들은 ML-IRL의 암시적 하이퍼그레이디언트를 위한 샘플 기반 추정치를 유도하였으며, 실현 가능한 내부 최적점에서 내부 헤시안이 궤적 피셔 정보 행렬에 비례함을 증명하였다. 이는 ML-IRL의 암시적 하이퍼그레이디언트 계산에 이 항등식을 적용한 첫 사례로 제시된다.
• 확장 가능한 솔버: 저자들은 정책-스코어 벡터를 스트리밍하는 SCFD 기반 솔버를 제안한다. 이 접근 방식은 스코어 행렬이나 밀집된 피셔 행렬을 생성하지 않으므로 메모리 요구 사항을 크게 줄인다.
• 경험적 검증: 제안된 방법은 이산 제어 환경(CartPole)과 연속 제어 환경(LQR)에서 1차 확률적 이중 레벨 베이스라인(single-loop ML-IRL) 및 명시적 피셔 솔버(explicit Fisher solver)와 비교 평가되었다.

결과
실험은 동일한 아키텍처와 하이퍼파라미터 튜닝을 적용하여 동일한 계산 예산(24시간 훈련 제한) 하에서 수행되었다.

• 정책 성능: 피셔 기반 방법들은 경쟁력 있는 정책 품질을 달성했다. CartPole의 경우, 모든 방법이 전문가 수준의 리턴(return)에 도달했다. LQR의 경우, single-loop ML-IRL이 가장 좋은 정책 음의 로그 가능도(NLL)와 환경 리턴을 기록했으나, "Sketching을 적용한 Fisher" 방법과의 차이는 미미했다.
• 보상 순위 결정 (Reward Ranking): 피셔 기반 방법들은 강력한 보상 순위 결정 품질을 보여주었다. LQR에서 "Sketching을 적용한 Fisher" 방법은 학습된 보상과 실제 보상의 궤적 리턴 간의 순위 상관계수(RankCorrelation)에서 0.972를 기록하며, single-loop 베이스라인과 명시적 피셔 솔버를 모두 능가했다.
• 효율성:
  • 메모리: 스케칭은 고차원인 LQR 설정에서 가장 큰 이점을 제공했으며, 명시적 피셔 솔버 대비 피크 메모리 사용량을 최대 1.31배 감소시켰다.
  • 속도: CartPole에서는 스케칭이 1.29배의 속도 향상을 보였다. LQR에서는 중간 규모의 스케치 크기(예: $m=32$)가 약간의 속도 향상(1.04배)을 제공했으나, 더 큰 스케치($m=256$)는 최종 손실 값은 더 낮았음에도 불구하고 스케칭 프로세스의 오버헤드로 인해 명시적 솔버보다 느려졌다.
  • 안정성: 저자들은 스케칭된 솔버가 명시적 공식보다 훨씬 약한 댐핑 파라미터($\lambda$) 하에서도 안정적으로 유지됨을 발견했으며, 이는 스케칭 과정이 암시적 정규화 역할을 할 수 있음을 시사한다.

의의 및 주장
본 논문은 역-KL 최적화 맥락에서 내부 헤시안과 피셔 정보 행렬 사이의 직접적인 연결 고리를 확립함으로써 이중 레벨 IRL의 확장성 병목 현상을 해결한다고 주장한다. 이러한 이론적 통찰을 스트리밍 매트릭스 스케칭과 결러하여, 저자들은 명시적인 헤시안이나 피셔 행렬 구축에 따른 이차 메모리 비용 없이 암시적 하이퍼그레이디언트를 계산할 수 있음을 입증하였다.

저자들은 피셔 기반 방법이 최종 정책 품질 측면에서는 single-loop ML-IRL과 경쟁할 만한 수준이지만, 고차원 문제에 대해 보상 순위 결정 성능과 계산 효율성(특히 메모리 감소) 면에서 뚜렷한 이점을 제공한다고 겸허히 결론지었다. 본 연구는 적절한 스케치 크기가 최적화 품질, 메모리 사용량, 실행 시간 사이의 최적의 절충점(trade-off)을 제공함을 시사한다.