Flow Q-Learning 논문 리뷰

Offline RL에서 흔히 지적되는 핵심 bottleneck은 dataset distribution과 실제 정책이 유도하는 분포 간의 support 불일치로 인해, 데이터셋에 없는 state-action pair에 대한 $Q(s,a)$의 잘못된 추정에 있다. 그로 인한 문제를 해결하는 몇 가지 대표적인 방법으로

• Conservative Q (CQL): 관측되지 않은 state-action pair에 대해 Q function 자체를 conservative하게 만들어(의도적으로 낮게 추정하여) 데이터셋에 없는 영역에 대해 과도한 낙관을 억제한다. 이는 정책이 데이터 support 밖으로 나가는 것을 간접적으로 방지하지만, 지나치게 보수적일 경우 성능 자체를 낮출 수 있다.
• Behavior Cloning: 학습된 정책이 데이터셋을 수집한 behavior policy에서 크게 벗어나지 않도록 제약을 가한다. KL divergence 등의 정규화를 통해 support mismatch를 줄일 수 있으나, 정책의 표현력이 제한되며 멀티모달 행동 분포를 다루는 데에는 한계가 있다.
• Stationary Distribution Ratio Correction(DICE): 목표 정책이 유도하는 stationary distribution과 데이터셋 분포 간의 비율을 직접 추정함으로써, 분포 불일치를 이론적으로 교정하려는 접근이다. 이는 강력한 수학적 기반을 가지지만, ratio 추정 자체가 어려우며 실제 구현에서는 안정성 문제가 자주 발생한다.

이러한 방법들은 서로 다른 관점에서 접근하지만, 결과적으로 데이터에 없는 행동을 어떻게 다룰지에 대해서 답하고자 한다. 그러나 본 논문에서는 이러한 벙법론들이 근본적으로 정책의 표현 형태, 즉 policy를 어떻게 parameterize하는지에 대해서는 비교적 보수적(ex: Gaussian)이라는 지적을 한다.

 

한편 최근 제안된 flow 및 diffusion 기반 정책은 복잡하고 멀티모달한 행동 분포를 직접 모델링할 수 있다는 가능성을 제시한다. 그러나 이러한 정책을 actor–critic 구조의 Offline RL에 직접 결합할 경우, RL 특유의 recursive한 bootstrap 구조와 iterative generative process가 맞물리며 단순한 불안정성을 넘어 심각한 학습 실패로 이어지는 문제가 발생한다. 따라서 본 논문은 Flow Q-Learning(FQL)을 제시하여, 새로운 접근으로 행동 분포의 표현력과 Offline RL의 안정성을 구조적으로 확보하고자 한다.

 

Preliminaries

우선 두 가지 질문에 대해서 답해보면서, 왜 FQL이 그러한 방법론에 도달하게 되었는지 알아보자.

왜 Gaussian Policy가 본질적인 병목이 되는가?

Offline RL에서 정책은 전통적으로 Gaussian 분포로 parameterize되어 왔다. 이는 reparameterization trick을 통해 안정적인 gradient 계산이 가능하며, actor-critic 구조와의 결합이 용이하다는 실용적인 강점 때문이다. 그러나 이러한 선택은 정책의 표현력 자체가 강하게 제한된다는 tradeoff가 존재할 수밖에 없어진다.

 

우선 Gaussian Policy는 기본적으로 unimodal이다. 이는 하나의 상태에 대해서 여러 개의 질적으로 다른 행동 모드가 공존하는 경우를 자연스럽게 표현하지 못한다는 문제가 있다. 이러한 상황에서 Gaussian policy는 여러 모드를 평균내어, 실제로는 데이터셋에 존재하지 않는 행동을 생성하거나, 하나의 모드만을 임의로 선택하게 된다.

 

이 문제는 Offline RL에서 특히 치명적인데, 이는 policy의 제한된 parameterization가 Q를 망가뜨릴 수 있는 가능성을 가지기 때문이다. Online RL에서는 환경 상호작용이 이를 교정할 수 있지만, Offline RL에서는 제한된 표현으로 인한 오류가 누적되어 위험성이 더 크다. 결국 기존의 Offline RL 방법들은, 다양한 제약식을 바탕으로 policy가 어긋나지 못하게 하는 방식으로 문제를 완화해왔다.

 

그러나 이는 근본적인 문제 해결이 아니며, 그러한 까닭이 본 논문에서 flow 기반 policy parameterization을 제시한 이유이기도 하다.

 

Flow policy가 그렇다면 왜 어려웠는가?

사실 Gaussian의 명백한 한계를 모를 사람들이 아니다. 그러면 단순하게 flow로 해도 되지 않느냐는 의문에 대해서, 왜 그것이 지금까지 어려웠는지에 대한 논의를 해보고자 한다.

우선 이상적인 경우에, flow나 diffusion 기반 정책은 gaussian 정책의 한계를 정면으로 돌파한다. 간단하게 flow 기반 정책에 대해서 이야기해보자. 우선 gaussian noise $Z \sim \mathcal N (0, I) $로 시작하여, 반복적인 변환을 통해 복잡하고 멀티모달한 행동 분포를 생성한다. 예를 들어서, flow policy는 다음과 같은 ODE로 정의된다:

$$ \frac{dx_t}{dt}=v_\theta (t,s,x_t), \qquad a=x_1 $$

이는 이론적으로 강력한 정책 표현이며, 특히 생성 모델에서 이미 그 강력함을 입증해왔다. 문제는 이를 RL loss에 연결하는 순간에 발생한다. Flow Matching(FM) 계열 중에서 Optimal Transport(OT)를 사용한 flow 모델의 loss를 우선 보면 다음과 같다:

학습의 불안정성이 크게 문제되지 않는 이유는, 이미 고정되어 있는 $x_0, x_1$을 참고하여, 특정 time $t$의 벡터장만을 학습하면 되기 때문이다. 강화학습에서 이러한 것이 어려운 이유는, actor가 critic을 참고하여 업데이트되고, critic은 actor을 참고하여 업데이트되는 recursive한 구조이기 때문이다. 이 구조 속에서 flow 기반 정책을 그대로 RL loss에 연결하면, 생성 모델에서 가정되던 안정적인 감독 조건이 즉시 붕괴한다. 이를 보다 명확히 이해하기 위해, 강화학습에서 actor–critic 업데이트가 어떤 형태의 수식을 갖는지 먼저 살펴보자.

 

우선 flow policy를 actor로 사용해보면, 노이즈 $z\sim \mathcal N (0,I)$에서 시작하여 ODE로 action을 생성해야 한다.

$$ \frac{dx_t}{dt}=v_\theta (t,s,x_t), \qquad a=\mu_\theta (s,z):=x_1 $$

Offline / actor-critic RL에서는 policy gradient는 보통 다음과 같다.

$$ J(\theta) = \mathbb E_{s\sim D, z\sim N}[Q_{\phi}(s, \mu_\theta(s,z)] \Rightarrow \nabla_\theta J(\theta) = \mathbb{E}_{s,z}\left[\nabla_a Q_\phi (s,a) \frac{\partial \mu_\theta (s,z)}{\partial\theta}\right]$$

정확히 $\frac{\partial\mu_\theta(s,a)}{\partial \theta} $가 문제가 되는데, 이는 $\mu_\theta$ 자체가 ODE의 결과이기 때문이다.

$$ \frac{\partial\mu_\theta}{\partial\theta}=\frac{\partial}{\partial\theta}\int_0^1 v_\theta (\tau, s, x_\tau) d\tau=\int_0^1 \frac{\partial v_\theta}{\partial \theta}+\frac{\partial v_\theta}{\partial x}\frac{\partial x_\tau}{\partial \theta}d\tau$$

이러한 형태는 1) Backpropagation Through Time(BPTT), 2) Jacobian 누적, 3) Gradient에서의 issue로 불안정하다는 문제가 있다.  이렇게 불안정한 gradient에 이어서, critic 또한 움직이는 target이라는 점에서 그러한 문제는 증폭될 수밖에 없는 구조가 나온다.

 

BPTT는 flow 모델의 전반에 존재하는 문제 아니냐는 지적은 충분히 가능하고, 그러한 말이 사실이다. 다만 생성 모델에서는 고정된 감독 신호를 회귀하거나 로그 우도를 최적화하는 구조이기 때문에, 이 문제가 실제 학습 불안정성으로 이어지지 않는다. 반면 강화학습에서는 ODE의 출력이 정책 행동으로 사용되고, 그 행동에 대한 가치 함수의 gradient가 다시 정책을 업데이트하는 자기참조적 구조를 형성한다. 이 경우 ODE의 시간적 재귀성과 actor–critic의 상호 의존성이 결합되어, gradient 폭발이나 불안정성이 구조적으로 발생한다.

Flow Q-Learning

항상 느끼지만 논문 저자 분께서 논문을 정말 보기 좋게, 친절하게 잘 쓰시는 것 같다.

이제 지금까지 flow 계열 policy가 어떤 문제를 안고 있었는지 대략적으로 확인했으니, 본 논문에서는 이러한 문제를 어떻게 해결했는지를 확인해보자. 위의 그림에서 나온 BC flow policy, One-step policy 중에서 우선 BC flow policy에 대해서 알아보도록 하자.

 

우선 BC flow policy이다. 핵심은 BC는 supervised learning이니, 얼마든지 flow기반 learning을 해도 문제가 없다는 것이다.

OT 기반으로 flow matching을 수행한다. 이로써 behavior policy를 간단하게 얻을 수 있다. 이제 One-step policy로 넘어가보자. One step policy는 기존의 expected reward maximization을 하면서, behavior policy에 distillation되도록 한다. 즉

위와 같이 distillation loss를 추가하여 Expected reward maximization을 하는 간단한 식이다. Distillation loss는 아래와 같다.

여기서 핵심은, 학습되는 실제 policy인 $\mu_\theta (s,z)$는 flow 계열 모델이 아니라는 점이다. 그러한 구조에 대해서 지적할 수도 있어서, 본 논문의 저자는 다음과 같은 정당화를 제공한다:

 

결론부터 말하자면, distillation loss는 Wasserstein Regularization과 연관되어 있다는 것이다. 단순히 다른 offline RL 알고리즘들에서 f-divergence 기반 behavior regularization을 적용하는 것이 아니라, Wasserstein distance를 사용하여 행동 공간의 거리 구조마저도 반영하는 metric-aware regularizer이기 때문에 policy family가 다르더라도 의미 있는 alignment가 가능하다는 것을 시사한다.

 

다른 방법론들과 비교하면서 왜 FQL이 더 나은지에 대해서 논문에서 친절히 설명하지만, 이 리뷰를 보는 독자들은 추가적으로 확인하길 바란다. (이 논문의 저자는 정말 논문을 잘 쓰는 것 같다는 생각이 매번 든다)

논문에서 제시하는 이 방법론의 강점 중에, 제일 인상깊었던 부분은 튜닝과 implementation이 정말 쉽다는 것이다. Flow를 RL에서 적용하기 위해서 추가적인 테크닉이라든지, 복잡한 고려 사항이 없다는 것이 이 논문을 매력적으로 만든다.

 

Experiments

https://seohong.me/projects/fql/

FQL

특히 잘 설명되어 있으니 논문 저자의 설명을 보는 것이 더 좋을 것이다.

가장 인상적인 점 중 하나는, Offline-to-Online RL에서 강점을 보인다는 것이고, 이를 위한 추가적인 아키텍처를 따로 도입하지 않아도 된다는 것이다.

 

Limitations

마지막으로 논문의 limitation에 대해서 알아보고 가자.

FQL은 flow policy의 표현력과 offline RL의 안정성을 효과적으로 결합했지만, 학습 중 ODE 계산 비용, 온라인 fine-tuning에서의 탐험 부재, 그리고 아직 실세계 로봇에서의 검증 부족이라는 한계를 가진다. 다만 이러한 한계들은 구조적인 결함이라기보다는, 향후 확장을 통해 충분히 보완 가능한 방향으로 남아 있다.

 

또 내가 생각하는 문제점은, FQL 역시 actor–critic 구조에서 완전히 자유롭지는 않다는 점이다.

특히 초기 critic의 불안정성, value maximization과 distillation 간의 긴장 관계, 그리고 BC flow policy 자체가 품질이 낮을 경우의 위험은 여전히 남아 있다. 또한 논문에서 은근하게 BC flow policy가 reasonable하다고 가정하는 느낌이 들어 다른 상황에서는 어떠한지에 대한 의문 또한 든다.

 

하지만 flow 기반으로 어떻게 offline learning을 해야 하는지 지침을 제공해주었다는 점에서 정말 좋은 논문이며 나도 비슷하게 아이디어를 가져보겠다는 생각이 많이 들게 되었다.