보상 인식 일관성 경로 증류로 오프라인 강화학습 가속

초록

본 논문은 사전 학습된 확산 플래너를 교사 모델로 활용하고, 독립적인 보상 모델을 결합한 새로운 일관성 경로 증류 기법을 제안한다. 보상 신호를 직접 손실에 포함시켜 단일 스텝 샘플링 학생 모델이 고보상 행동 모드를 선택하도록 유도함으로써, 오프라인 강화학습에서 서브옵티멀 데이터에도 강인한 성능을 달성한다. 실험 결과는 MuJoCo, FrankaKitchen, Maze2D 등에서 기존 최첨단 대비 9.7% 향상과 142배 이상의 추론 속도 향상을 보여준다.

상세 분석

이 연구는 확산 모델이 복잡한 의사결정 문제에서 뛰어난 표현력을 갖지만, 다수의 디노이징 스텝이 필요해 실시간 적용에 한계가 있다는 점을 출발점으로 삼는다. 기존의 일관성 모델(Consistency Model) 기반 가속 방법은 이미지 생성 분야에서 성공을 거두었으나, 강화학습에 적용될 때 두 가지 주요 문제에 직면한다. 첫째, 행동 복제(Behavior Cloning) 방식은 전문가 수준 데이터에만 효과적이며, 중간 품질 혹은 다양한 정책이 섞인 오프라인 데이터에서는 다중 모드 중 저보상 모드까지 학습하게 된다. 둘째, 액터-크리틱 구조를 차용한 방법은 교사와 학생, 그리고 가치 네트워크를 동시에 학습해야 하므로 하이퍼파라미터 튜닝과 학습 안정성 확보에 큰 비용이 든다.

논문은 이러한 한계를 극복하기 위해 “보상 인식 일관성 경로 증류(Reward‑Aware Consistency Trajectory Distillation, RA‑CTD)”라는 프레임워크를 설계한다. 핵심 아이디어는 (1) 사전 학습된 확산 플래너를 교사 모델로 고정하고, (2) 별도로 학습된 반환‑투‑고(return‑to‑go) 보상 모델을 사용해 학생 모델의 출력에 직접 보상 손실을 부과한다는 것이다. 여기서 보상 손실 L_reward = −Rψ(s_n, â_n) 은 학생이 생성한 첫 번째 행동 â_n 에 대한 예상 누적 보상을 최대화하도록 유도한다. 이 손실은 기존의 CTM 손실(L_CTM)과 DSM 손실(L_DSM)과 가중치 α, β, σ 로 결합되어 전체 목표 L = αL_CTM + βL_DSM + σL_reward 를 형성한다.

CTM 손실은 임의의 두 시점 k<u<t 사이의 경로 일관성을 강제함으로써 학생이 “anytime‑to‑anytime” 점프를 학습하게 하고, DSM 손실은 청정 데이터 공간에서의 복원 정확도를 유지한다. 보상 손실을 추가함으로써 학생은 다중 모드 분포 중 고보상 모드에 집중하도록 학습되며, 이는 기존의 행동 복제 방식이 저보상 모드까지 학습하는 문제를 자연스럽게 해결한다. 또한, 보상 모델이 노이즈가 없는 청정 행동 공간에서 학습되므로, 노이즈‑인식 보상 모델을 별도로 설계할 필요가 없어 구현 복잡도가 크게 낮아진다.

학습 과정은 완전히 디커플링되어 있다. 교사 확산 플래너는 기존 EDM(Elucidated Diffusion Model) 방식으로 사전 학습되고, 보상 모델은 오프라인 데이터에서 반환‑투‑고를 예측하도록 별도 학습된다. 이후 학생 모델은 교사와 보상 모델을 고정한 상태에서 위 손실들을 동시에 최소화한다. 이 구조는 (i) 교사‑학생 간의 파라미터 공유가 없으므로 학습 안정성이 높고, (ii) 하이퍼파라미터 탐색 범위가 크게 축소되며, (iii) 멀티‑네트워크 동시 최적화에 따른 메모리·연산 부담이 사라진다.

실험에서는 D4RL 기반 MuJoCo 연속 제어, FrankaKitchen 로봇 조작, 그리고 Maze2D 장기 계획 과제를 사용하였다. 결과는 RA‑CTD가 기존 최첨단(예: Diffusion‑based Planner, Consistency Trajectory Models, Actor‑Critic 기반 확산) 대비 평균 9.7%의 누적 보상 향상을 보였으며, 추론 단계에서는 단일 스텝 샘플링 덕분에 142배 이상의 속도 개선을 달성했다. 특히 서브옵티멀 데이터가 다수 포함된 “medium‑replay” 설정에서도 보상 인식 손실이 효과적으로 고보상 모드를 선택함을 확인하였다.

이 논문의 주요 기여는 (1) 보상 신호를 일관성 증류 과정에 직접 통합한 새로운 학습 목표 설계, (2) 교사·학생·보상 모델을 완전히 독립적으로 학습함으로써 훈련 복잡성을 크게 낮춘 시스템 아키텍처, (3) 실험을 통해 단일 스텝 샘플링이 고성능을 유지하면서도 실시간 적용 가능성을 입증한 점이다. 향후 연구에서는 보상 모델의 불확실성을 고려한 베이지안식 보정, 멀티‑에이전트 협업 시나리오, 그리고 비정형 관측(이미지, 라이다) 기반 플래너 확장 등이 기대된다.