Nonparametric Bayesian Inverse Reinforcement Learning with Data-Parallel Gibbs Sampling

기술 요약: 데이터 병렬 깁스 샘플링을 이용한 비매개변수적 베이지안 역강화학습

문제 정의

역강화학습(Inverse Reinforcement Learning, IRL)은 전문가의 시연으로부터 보상 함수를 복구하는 것을 목표로 한다. 그러나 일반적인 정식화는 대개 단일한 동질적 전문가를 가정한다. 서로 다른 선호도를 가진 여러 전문가로부터 수집된 데이터를 통합할 경우, 매개변수적 방법들은 종종 평균화된 보상을 복구하게 되어 어떤 개별 전문가에게도 부합하지 않는 결과를 초래한다. 또한, 기존의 베이지안 접근 방식은 보상 유형(클러스터)의 수를 사전에 고정해야 하는 경우가 많은데, 이는 잠재적 구조를 발견하고자 하는 목적에서 불합리한 가정이다. 아울러, 보상 사후 확률에 대한 마르코프 체인 몬테카를로(MCMC) 추론은 매 샘플마다 내부 계획(가치 반복)이 필요하기 때문에 계산 비용이 매우 높으며, 이는 비매개변수 모델에서 추론 비용이 궤적 수와 클러스터 수 모두에 따라 증가하는 병목 현상을 야기한다.

방법론

저자들은 비매개변수적 베이지안 IRL 프레임워크와 데이터 병렬 깁스 샘플링을 결합하여 이러한 과제들을 해결하는 시스템을 제안한다.

• 생성 모델: 시스템은 보상 함수에 디리클레 프로세스(Dirichlet Process, DP) 사전 확률을 배치한다. 이를 통해 클러스터 수를 사전에 고정하는 대신, 보상 자체와 함께 잠재적 보상 유형의 수를 함께 추론할 수 있다. 보상 함수는 $R(s) = \phi(s)^\top w$로 매개변수화되며, 여기서 $\phi(s)$는 특징 벡터이고 $w$는 가중치 벡터이다. 가능도(Likelihood)는 소프트 가치 반복(soft value iteration)에 의해 결정되는 볼츠만 합리적(Boltzmann-rational) 전문가 모델을 따른다.
• 추론 알고리즘: 저자들은 두 가지 업데이트를 한 번의 스윕(sweep) 동안 교대로 수행하는 **축약된 깁스 샘플러(collapsed Gibbs sampler)**를 구현한다.
  1. 클러스터 할당 업데이트: 집중 매개변수 $\alpha$를 기반으로 궤적을 기존 클러스터에 할당하거나 새로운 클러스터에 할당하기 위해 중국인 식당 프로세스(Chinese Restaurant Process, CRP) 업데이트를 사용한다.
  2. 가중치 업데이트: 새로운 보상 가중치를 제안하기 위해 메트로폴리스-헤이스팅스(Metropolis-Hastings, MH) 업데이트를 사용한다.
  3. 계획(Planning): 가능도 평가를 위한 내부 계획 루틴으로 100회의 고정된 벨만 백업(Bellman backups)을 포함하는 소프트 가치 반복을 사용한다.
• 병렬화: 계산 비용을 해결하기 위해, HPC 하드웨어 상에서 Ray를 사용하여 CPU 코어에 걸쳐 샘플러를 병렬화하였다. 궤적은 워커(worker)들에 분할되어 각각 독립적인 깁스 스윕을 실행한다. 컨센서스 머지(consensus merge) 단계는 각 스윕 후에 결과를 집계하며, 대응하는 클러스터를 식별하기 위해 허용 오차($10^{-6}$)를 가진 요소별 비교를 사용한다.
• 구현: 시스템은 약 2,100라인의 파이썬 코드로 구성된 JAX로 구축되었으며, 클러스터 재라벨링(relabeling)의 유연성을 유지하면서 볼츠만 로그 정책 계산을 위한 JIT 컴파일을 활용하는 동시에 외부 제어 흐름을 위한 일반 파이썬을 유지한다.

주요 기여

1. 완전한 구현: CRP 및 MH 업데이트를 사용하는 축약된 깁스 샘플러를 갖춘 비매개변수적 베이지안 IRL의 오픈 소스 구현을 제공하며, 이를 ObjectWorld 벤치마크의 그라운드 트루스 보상에 대해 검증하였다.
2. 병렬 추론 엔진: Ray를 사용한 데이터 병렬 확장형을 통해 상당한 속도 향상을 달eric하고, 컨센서스 머지 휴리스틱에서 발생하는 처리량과 정확도 사이의 트레이드오프를 규명하였다.
3. 경험적 평가: 두 가지 및 세 가지 전문가 유형에 대한 엄격한 평가를 수행하였으며, 환경의 특징 구별 가능성(feature discriminability)이 샘플러의 동작과 독립적으로 어떻게 회복 품질을 제한하는지에 대한 분석을 포함하였다.
4. 재현성: HPC 인프라에서 직접 복제가 가능하도록 컨테이너화된 시스템(Docker 이미지 및 SLURM 스크립트)을 제공한다.

결과

시스템은 객체의 색상과 거리를 인코딩하는 이진 특징 벡터를 사용하는 $10 \times 10$ ObjectWorld 그리드에서 평가되었다.

• 두 전문가 설정 ($K=2$): 직렬 샘플러는 수정 랜드 지수(Adjusted Rand Index, ARI) 1.000으로 $K=2$를 성공적으로 복구하였으며, 이는 단일한 평균 보상으로 붕괴되는 최대 엔트로피(MaxEnt) IRL 베이스라인(ARI=0.000)을 크게 상회하는 결과이다. 시스템은 약 134초의 벽 시간(wall time) 동안 500회의 스윕 내에 이를 달성하였다.
• 세 전문가 설정 ($K=3$): 샘플러는 모든 실행에서 클러스터 수($K=3$)를 올바르게 식별하였다. 그러나 할당 ARI는 0.48에서 0.58 사이였다. 사후 분석 결과, 이러한 한계는 샘플러나 하이퍼파라미터 때문이 아니라 특정 환경 인스턴스에 의한 행동적 중첩(behavioral overlap) 때문임이 밝혀졌다. 많은 랜덤 시드에서 특정 객체 색상이 부재하였고, 이로 인해 서로 다른 보상 유형(예: "파란색 애호가"와 "검은색 애호가")이 구별 불가능한 궤적 분포를 생성하였다. 저자들은 신뢰할 수 있는 $K=3$ 평가를 위해서는 랜덤 시딩이 아닌 통제된 객체 배치가 필요하다고 결론지었다.
• 병렬 속도 향상: HPC 하드웨어(AMD EPYC 노드)에서 시스템은 8개의 워커를 사용하여 **4.79×**의 피크 속도 향상을 달성하였다.
• 처리량-정확도 트레이드오프: 더 많은 워커를 사용할수록 벽 시간은 감소했지만, ARI는 급격히 저하되었다(16개 워커에서 0.035로 하락). 이는 컨센서스 머지 단계 때문인 것으로 분석되었다. 독립적인 MH 업데이트가 가중치 벡터를 미세하게 드리프트(drift)시켰고, 머지 휴리스틱이 이 드리프트가 허용 오차($10^{-6}$)를 초과했을 때 서로 다른 클러스터를 잘못 식별하였다. $\ell_2$ 가중치 오차는 안정적으로 유지되었는데, 이는 워커별 보상 추론은 정확했으나 집계 과정에서 클러스터 정체성이 상실되었음을 나타낸다.

의의 및 주장

본 논문은 비매개변수적 베이지안 IRL이 클러스터 수를 사전에 지정하지 않고도 전문가 유형의 수와 특정 보상 함수를 효과적으로 복구할 수 있음을 입증한다고 주장한다. 시스템은 환경이 충분한 특징 구별 가능성을 제공하는 시나리오($K=2$ 케이스)에서 완벽한 복구를 달성한다.

저자들은 $K=3$ 사례에 관한 연구 결과를 알고리즘의 실패가 아닌, 합성 IRL 벤치마크 설계에 대한 중요한 통찰로 겸허하게 제시하였다. 즉, 특징 구별 가능성은 보상 명세로부터 당연히 가정될 수 없으며, 특정 환경 인스턴스에 대해 반드시 검증되어야 한다는 것이다.

병렬화와 관련하여, 논문은 컨센서스 머지 단계를 처리량-정확도 트레이드오프의 원인으로 지목하였다. 저자들은 이것이 알고리즘의 근본적인 한계라기보다 엔지니어링 이슈라고 주장하며, 중앙 집중식 가중치 업데이트를 통해 처리량 이득을 유지하면서 드리프트 문제를 해결할 수 있을 것이라고 제안한다. 본 연구는 확장 가능한 비매개변수적 IRL의 토대를 마련하였으나, 머지 메커니즘을 개선하고 더 다양한 절차적 생성 환경에서 테스트하기 위한 후속 연구가 필요함을 인정하고 있다.