한 단계 흐름 매칭으로 최대 엔트로피 강화학습 가속화

초록

FLAME은 흐름 매칭(Flow Matching) 기반 정책을 최대 엔트로피 강화학습에 직접 통합한다. Q‑값을 이용한 중요도 재가중으로 에너지 기반 목표의 정규화 상수를 제거하고, 편향을 보정한 분리형 엔트로피 추정기로 탐색 효율을 유지한다. MeanFlow 기법을 적용해 단일 ODE 단계(NFE=1)로 고품질 행동을 생성하며, MuJoCo 실험에서 기존 가우시안 정책을 능가하고 다단계 확산 정책과 동등한 성능을 보이면서 추론 비용을 크게 낮춘다.

상세 분석

본 논문은 연속 제어에서 표현력이 풍부한 생성 모델을 활용하려는 최근 흐름‑기반 정책 연구와, 탐색‑활용 균형을 정식화한 최대 엔트로피 강화학습(MaxEnt RL) 사이의 격차를 메우는 데 초점을 맞춘다. 기존 확산 정책은 다수의 함수 평가(NFE)가 필요해 실시간 제어에 부적합하고, 흐름 매칭(Flow Matching, FM)은 목표 샘플이 필요해 에너지 기반(MaxEnt) 목표와 직접 결합하기 어렵다. 논문은 세 가지 핵심 기여를 제시한다.

첫째, Q‑Reweighted Flow Matching(QRFM) 목표를 도입한다. MaxEnt 정책의 에너지 형태 π∗(a|s)∝exp(Q(s,a)/α)에서 정규화 상수 Z(s)가 불가능한데, 저자는 임의의 양의 가중치 함수 g(a,s)를 곱해 흐름 매칭 손실을 재가중함으로써 Z(s)를 완전히 소거한다. 구체적으로, g_maxent(a_t,s)=h_t(a_t|s)·Z(s)·p_t(a_t|s) 로 설정하고, 역샘플링 트릭을 이용해 a_t를 제안 분포 h_t에서 샘플링한 뒤, 정규분포 ϕ_{1|t}(a_1|a_t) 로 a_1을 역생성한다. 이렇게 하면 기대값 안에 exp(Q(s,a_1)/α)만 남아, Q‑값만으로 정책을 학습할 수 있다. 이 과정은 조건부 흐름 매칭(CFM)과 동일한 최적 벡터 필드를 보장하므로, 기존 FM 이론과 완벽히 호환된다.

둘째, 연속 흐름의 로그가능도 추정에 내재된 이산화 편향을 정량화하고, 이를 보정하는 두 가지 전략을 제시한다. FLAME‑R은 증강 ODE(augmented‑ODE) 형태로 로그밀도 변화를 정확히 적분해 무편향 추정을 제공한다. 반면 FLAME‑M은 다단계 적분을 이용한 분리형 엔트로피 추정기로, 학습 단계에서는 정확한 로그가능도를 사용해 엔트로피 보너스를 계산하고, 실행 단계에서는 단일 ODE 스텝만 수행한다. 이렇게 하면 학습 시 안정적인 탐색을 유지하면서도 실시간 추론 속도를 크게 향상시킬 수 있다.

셋째, MeanFlow 프레임워크를 결합해 벡터 필드의 평균 속도를 학습한다. 기존 FM은 경로가 곡선이면 다중 스텝 통합이 필요했지만, MeanFlow는 평균 속도를 직접 추정해 경로 곡률을 최소화한다. 결과적으로 NFE=1인 한 단계 흐름으로도 고품질 다중모달 행동을 생성할 수 있다.

실험에서는 MuJoCo 연속 제어 벤치마크에서 FLAME‑M과 FLAME‑R이 모두 Gaussian 기반 SAC보다 높은 평균 보상을 기록했으며, 4~8배 적은 NFE를 요구하는 다단계 확산 정책과 거의 동등한 성능을 보였다. 특히 복잡한 다중모달 행동이 요구되는 환경에서 FLAME은 모드 붕괴 없이 안정적인 탐색을 유지한다.

전체적으로 논문은 (1) 정규화 상수 제거를 통한 Q‑값 기반 흐름 매칭, (2) 편향 보정된 엔트로피 추정, (3) MeanFlow 기반 단일 스텝 제어라는 세 축을 통해 MaxEnt RL과 흐름 기반 생성 모델을 성공적으로 융합했다. 다만, 현재 구현은 Gaussian 베이스 p_0와 선형 OT 커플링에 의존하므로, 고차원 복잡한 행동 공간에서의 확장성 및 비선형 커플링 적용에 대한 추가 연구가 필요하다.