PolicyFlow: Policy Optimization with Continuous Normalizing Flow in Reinforcement Learning

기술 요약: PolicyFlow

문제 정의

근사 정책 최적화(Proximal Policy Optimization, PPO)는 안정성과 단순성으로 잘 알려진 표준 온폴리시(on-policy) 강화 학습(RL) 알고리즘입니다. 그러나 PPO는 일반적으로 가우시안 정책 분포에 의존하며, 이는 복잡하거나 다봉형(multimodal) 또는 심하게 왜곡된 액션 분포를 표현하는 데 어려움을 겪습니다. 연속 정규화 흐름(Continuous Normalizing Flows, CNF) 및 확산 모델(diffusion models)과 같은 표현력이 풍부한 생성 모델은 이러한 복잡한 분포를 포착할 수 있지만, 이를 PPO에 통합하는 것은 상당한 과제를 안겨줍니다.

주요 병목 현상은 PPO 서로게이트 목적 함수를 계산하기 위해 중요도 비율($\pi(a|s)/\hat{\pi}(a|s)$)을 계산하는 과정에서 발생하는 우도 평가(likelihood evaluation) 요구사항입니다. CNF의 경우, 정확한 우도 평가를 위해서는 전체 상미분 방정식(ODE) 궤적을 시뮬레이션하고 경로별 역전파(path-wise backpropagation)를 수행해야 합니다. 이 과정은 계산 비용이 많이 들고 메모리 집약적이며, 그래디언트 폭주 또는 소실로 인해 수치적으로 불안정할 수 있습니다. 기존의 PPO 스타일 업데이트를 흐름 기반 정책에 적용하려는 시도들(예: FPO, DPPO)은 비대칭적 추정 편향(asymmetric estimation bias)을 겪거나, 오프 매니폴드(off-manifold) 탐색의 어려움으로 인해 처음부터 학습할 때 성능이 저하되는 문제를 보였습니다. 또한, 표준 엔트로피 정규화 기법은 흐름 기반 정책에 적용하기 어려운데, 이는 정책 엔트로피를 명시적으로 계산하기 위해 값비싼 발산 적분(divergence integration)이 필요하기 때문입니다.

방법론

저자들은 전체 궤적 우도 평가를 요구하지 않으면서도 표현력이 풍부한 CNF 정책을 PPO 스타일의 목적 함수와 통합하는 온폴리시 RL 알고리즘인 PolicyFlow를 제안합니다. 이 방법은 두 가지 핵심 구성 요소로 이루어져 있습니다.

1. 속도장 변동을 통한 중요도 비율 근사

현재 정책 $\phi_1$과 참조 정책 $\hat{\phi}_1$ 사이의 최종 변화량(terminal shift)을 계산하기 위해 전체 ODE 궤적을 시뮬레이션하는 대신, PolicyFlow는 단순한 선형 보간 경로를 따라 **속도장 변동(velocity field variations)**을 사용하여 중요도 비율을 근사합니다.

• 정책 표현: 정책은 신경망 속도장 $v_t$에 의해 제어되는 흐름 $\phi_t(z; s)$를 적분하고, 최종 시간에 가우시안 노이즈를 추가하여 액션을 생성합니다: $a = \phi_1(z; s) + n$.
• 근사: 최종 변화량 $\delta\phi_1 = \phi_1 - \hat{\phi}_1$은 선형 보간 경로 $x_t = (1-t)z + t\hat{\phi}_1(z; s)$를 따라 속도장의 차이 $\delta v_t = v_t - \hat{v}_t$를 적분함으로써 근사됩니다.
• 서로게이트 목적 함수: 중요도 비율 $\rho$는 전체 흐름 궤적이 아닌 속도장 변동의 가우시안 우도 비율을 사용하여 추정됩니다. 이를 통해 업데이트 단계에서 전체 ODE 시뮬레이션을 통한 역전파 없이도 보간 경로를 따라 효율적으로 훈련 목적 함수를 추정할 수 있습니다. 이론적 분석에 따르면, 이는 로그 우도에 1차 오차만을 도입하며, 이는 PPO의 클리핑 메커니즘에 의해 자연스럽게 제한됩니다.

2. 엔트로피를 위한 브라운 운동 정규화기(Brownian Regularizer)

모드 붕괴(mode collapse)를 방기하고 다양한 탐색을 장려하기 위해, 흐름의 난해한 엔트로피를 명시적으로 계산하는 대신 저자들은 브라운 운동 정규화기를 도입합니다.

• 개념: 브라운 운동과 열 방정식에서 영감을 얻은 이 정규화기는 학습된 속도장이 엔트로피가 증가하는 역학을 따르도록 유도합니다.
• 메커니즘: 이는 rectified flows에서의 속도장과 스코어 함수($\nabla \log p_t$) 사이의 관계를 활용합니다. 이 정규화기는 학습된 속도장이 참조 흐름의 스코어로부터 유도된 타겟 벡터장으로부터 벗어나는 것에 대해 페널티를 부여합니다.
• 정식화: 정규화 항 $J_{Reg}$은 학습된 속도장과 타겟 정렬 항(참조 속도장 및 상태로부터 유도됨) 사이의 차이를 최소화하는 항과, 주입된 가우시안 노이즈의 분산을 장려하는 항을 포함합니다. 이 접근 방식은 값비싼 로그 우도 계산이나 휴리스틱한 노이즈 주입을 피하면서, 궤적 다양성을 유지하기 위한 경량화되고 원칙적인 방법을 제공합니다.

주요 기여

1. 효율적인 CNF-PPO 통합: PolicyFlow는 속도장 변동을 통해 중요도 비율을 근사함으로써, 전체 ODE 궤적에 대한 값비싼 경로별 역전파 없이도 표현력이 풍부한 CNF 정책을 온폴리시 RL에서 사용할 수 있게 합니다.
2. 브라운 운동 정규화기: 명시적인 엔트로피 계산 없이도 모드 붕괴를 완화하고 다양한 행동을 장려하는, 브라운 운동에서 영감을 받은 새롭고 계산 효율적인 엔트로피 정규화기를 제안합니다.
3. 실험적 검증: 이 방법은 MultiGoal, MuJoCo Playground, IsaacLab 벤치마크를 통해 검증되었으며, 가우시안 PPO 및 최신 흐름 기반 베이스라인(FPO, DPPO) 대비 경쟁력 있거나 우수한 성능을 입증했습니다.

실험 결과

저자들은 MultiGoal, MuJoCo Playground, IsaacLab 벤치마크에서 PolicyFlow를 평가했습니다.

• MultiGoal (모드 붕와 테스트): 여섯 개의 대칭적인 목표가 있는 환경에서 표준 PPO(가우시안)는 모든 목표에 도달하는 데 실패했습니다. 베이스라인인 FPO와 DPPO는 효과적인 엔트로피 정규화의 부족으로 인해 몇 개의 모드로 붕괴되었습니다. 브라운 운동 정규화기를 사용한 PolicyFlow는 균형 잡힌 다봉형 분포를 성공적으로 포착하여, 거의 동일한 확률로 모든 목표에 도달했습니다.
• MuJoCo Playground: PolicyFlow는 다양한 작업(예: WalkerWalk, CheetahRun)에서 FPO 및 DPPO와 대등하거나 이를 상회하는 성능을 달성했습니다. 종종 더 빠르게 수렴하여 높은 샘플 효율성을 보여주었습니다.
• IsaacLab: 로보틱스 작업(보행, 조작)에 대해 PPO와 비교되었습니다. PolicyFlow는 PPO의 점근적 성능(asymptotic performance)과 일치하거나 이를 능가하는 모습을 지속적으로 보여주었습니다.
• 계산 효율성: PolicyFlow는 샘플링을 위해 ODE 시뮬레이션을 포함하지만, (전체 ODE 역전파를 피하는) 훈련 단계는 표준 작업에 대한 PPO 대비 반복당 훈련 시간을 50% 미만으로 증가시켰으며, 임베딩 차원이 증가하더라도 2배 미만의 비용을 유지했습니다.

의의 및 주장

본 논문은 PolicyFlow가 온폴리시 RL을 위한 실용적이고 표현력이 풍부한 프레임워크를 제공한다고 주장합니다. 전체 궤적 우도 평가의 필요성을 분리함으로써, 연속 정규화 흐름의 표현 능력과 PPO의 안정성 사이의 간극을 메웁니다.

저자들은 다음을 강조합니다:

• 중요도 비율의 근사는 이론적으로 제한되어 있으며, PPO 클리핑 영역 내에서 안정적인 훈련에 충분합니다.
• 브라운 운동 정규화기는 탐색 유지를 위한 휴리스틱한 노이즈 주입이나 값비싼 발산 적분의 원칙적인 대안을 제공합니다.
• PolicyFlow는 최적화 안정성, 계산 비용 및 편향과 관련된 기존 흐름 기반 RL 방법(FPO, DPPO)의 특정 한계를 해결하여, 표현력이 풍부한 생성 정책을 온라인 강화 학습에서 더 쉽게 사용할 수 있도록 만듭니다.

이 연구는 PolicyFlow가 생성 모델링과 강화 학습을 잇는 토대 역할을 할 수 있음을 시사하며, 향-매칭(flow-matching) 정책의 미세 조정, 오프라인 RL로의 확장, 그리고 확산 모델과의 연결성 탐구와 같은 잠재적인 미래 방향을 제시합니다.