ReinFlow: Fine-tuning Flow Matching Policy with Online Reinforcement Learning 논문 리뷰

Summary

Robot manipulation task에서 모방 사전 학습을 한 RL으로 online에서 finetune하려는 시도는 diffusion policy에서 먼저 활발히 연구되었다. Diffusion은 denosing step 자체가 가우시안 전이로 정의되기 때문에 log-probability를 도출하기가 비교적 자연스러우며, 따라서 DPPO와 같이 policy gradient를 통한 online fine-tuning이 구현 가능한 형태로 나온다.

 

한편 flow matching 기반 정책은 sampling이 ODE로 정의되는 결정론적 경로에 가깝고, 이러한 이유로 prob 계산을 도출하는 것이 어렵다. 그러한 이유로 PPO류의 policy gradient 방법을 사용하는 것이 아닌 policy의 action에 대한 미분을 적용하는 방안으로 접근할 수밖에 없는데, 그렇게 되는 경우 BPTT(Backpropagation Through Time)이라는 문제점에 도달하게 된다. 이에 대한 대안으로 one-step policy distillation이 있지만, 결국 flow 모델의 표현력을 제한할 수밖에 없는 tradeoff의 결과에 도달하게 된다.

 

Reinflow는 flow에서 접근하지만, 결정론적 ODE에 Gaussian noise injection을 부여함으로써 flow를 discrete-time Markov process로 재정의하고, 그 joint likelihood를 닫힌 형태로 정확히 계산해내면서 PPO류의 policy gradient 방법을 적용하는 대안을 제시한다. 우선 이 논문을 이해하기 위한 preliminaries에 대해서 살펴보고 이후 구체적인 방법에 대해서 알아보도록 하자.

 

Preliminaries

이 논문을 이해할 때 가장 먼저 잡아야 하는 뼈대는 “정책경사 기반 RL이 연속 제어에서 왜 log-prob를 필요로 하느냐”이다. Score 기반 PPO/REINFORCE류의 기본형은 대개 다음 구조다.

$$\nabla_\theta J(\pi_\theta) = \mathbb{E}\big[ A^\pi(s,a)\,\nabla_\theta \log \pi_\theta(a|s)\big]$$

여기서 중요한 건, 정책 업데이트는 “action을 θ로 미분하는 pathwise gradient”가 아니라 “log-prob의 gradient”로 이루어진다는 점이다. 이 구조는 deterministic policy gradient에 기반한

$$\nabla_\theta J \approx \mathbb{E}\big[ \nabla_a Q(s,a)\,\frac{\partial a}{\partial \theta}\big]$$

 

같은 형태와 다르다. 그리고 둘의 차이가 바로 Summary에서 얘기해온 BPTT 폭발/소실 문제와 직결된다. 두 번째 식은 a가 다단계 생성 절차(denoising steps)를 거쳐 나오면 action에 대한 미분이 긴 체인룰 곱으로 누적되어 수치적으로 박살나기 쉽다. 반면 첫 번째 식은 action 생성 경로를 θ로 직접 미분하지 않고 “그 action이 나올 확률의 민감도”로 업데이트한다는 점에서, 적어도 그 형태의 체인 곱 문제를 피한다. 아래 그림을 참고해보면 좋다:

그러면 $\log \pi(a|s)$의 형태로 접근한다면, 그 값은 어떻게 계산할까? CNF에서는 이론적으로 정의된 식이 존재한다:

$$\log p_1(\psi_1(x)) = \log p_0(\psi_0(x)) - \int_0^1 \nabla\cdot v(t,\psi_t(x))\,dt$$

 

이렇게 적분 형태로 나오는데, 이것도 현실적으로 문제는 많다. 우선 이 적분을 수치적으로 풀어야 하고, divergence term은 Hutchinson estimator같은 것으로 trace를 추정해야 하는데, 문제는 두 가지다.

1. Monte-Carlo 오차: trace estimator는 표본 오차가 들어간다.
2. Discretization 오차: 적분을 K step으로 쪼개면 step 크기 $\Delta t$가 커질수록 오차가 커진다.

로봇 제어는 빠른 추론이 중요해서 “few-step이 좋은데”, CNF의 log-prob 근사는 precisely 그 few-step에서 가장 불안정해진다. 이 부분이 ReinFlow가 해결하려는 "수학적으로 문제가 되는 지점"이다.

 

Method

ReinFlow의 방법은 크게 세 덩어리로 보면 깔끔하다. (1) likelihood를 닫힌형으로 만들기, (2) 그걸 써서 policy gradient 정리를 세우기, (3) noise net과 정규화로 학습을 안정화하기.

(1) Likelihood Computation

Euler method의 deterministic map은 다음과 같다:

$$a_{k+1} = a_k + v_\theta(t_k, a_k, o)\Delta t$$

본 논문에서는 이러한 결정론적 전이를 사용하지 않고 전이를 가우시안으로 정의한다.

$$a_{k+1} \sim \mathcal{N}\Big(a_k + v_\theta(t_k, a_k, o)\Delta t, \ \sigma^2_{\theta'}(t_k, a_k, o)\Big)$$

여기서 $\sigma_{\theta'}$가 noise injection network다.

중요한 건 “노이즈를 넣었다”가 아니라 “이 노이즈를 모델 내부에 명시적으로 넣어서 전이확률을 닫힌형으로 정의했다”는 점이다. 그 결과 확률밀도를 도출할 수 있게 되는데, 그 확률밀도는 다음과 같이 정의된다:

$$\log \pi(a_0, \dots, a_K | o) = \log \mathcal{N}(a_0; 0, I) + \sum_{k=0}^{K-1} \log \mathcal{N}(a_{k+1}; \mu_k, \sigma_k^2)$$

 

이 표현은 CNF divergence 적분처럼 근사하는 게 아니다. “모델이 그렇게 정의되어 있기 때문에” 그대로 정확하다. 논문이 강조하는 ‘few-step에서도 exact’는 사실 여기서 나온다. CNF식은 step이 커지면 근사오차가 커지지만, ReinFlow의 log-prob는 step이 커져도 여전히 정확한 가우시안 로그우도들의 합이다(모델 정의가 discrete-time이니까).

여기서 subtle한 점이 하나 있다. ReinFlow가 사실상 하는 일은 “원래의 flow(ODE)를 그대로 쓰면서 log-prob만 계산한 것”이 아니다. 정책을 바꾼 것이다. 학습 중에는 deterministic flow가 아니라 noise-injected stochastic process로 정책을 재정의한다. 다만 drift는 pre-trained flow의 velocity field가 담당하고, stochasticity는 noise net이 담당하는 구조다.

(2) Marginal이 아닌 joint에 대한 Policy Gradient?

1번을 이해했다면 우리가 평소에 알고 있는 policy gradient를 적용할 때의 score가 아님을 알 수 있다. 우리는 최종 action에 대한 확률만을 필요로 하지만 논문에서 도출되는 결과는 그 action을 샘플링하는 데에 있어서 나오는 trajectory에 대한 joint 확률분포이기 때문이다. 논문은 이에 대해 새로운 PG theorem을 세운다.

핵심 형태는

$$\nabla_\theta J = \mathbb{E} \left[ A(o, a) \nabla_\theta \log \pi(a_0, \dots, a_K | o) \right]$$

그리고 Markov property 때문에 전체 궤적의 결합 확률은 초기 상태 확률과 각 스텝별 전이 확률의 곱으로 표현되며, 로그를 취하면 아래와 같이 합의 형태로 분해된다:

$$\log \pi(a_0, \dots, a_K | o) = \log p(a_0) + \sum_{k=0}^{K-1} \log \pi(a_{k+1} | a_k, o)$$

최종적인 업데이트는 다음과 같다:

$$\nabla_\theta J = \mathbb{E} \left[ A(o, a) \nabla_\theta \sum_{k=0}^{K-1} \log \pi(a_{k+1} | a_k, o) \right]$$

최종 action의 log-probability를 직접 게산하지 않고도, 디노이징 과정에서 joint prob의 gradient로 정책을 업데이트할 수 있다.

논문은 이를 기반으로 PPO clipped surrogate를 사용해 안정적으로 fine-tuning한다.

(3) Noise Injection Network와 Regularization

Noise net $\theta'$는 단순히 “탐험을 위한 노이즈 크기”가 아니라 정책의 일부로 학습된다. 학습 중에는 noise net이 샘플 경로에 다양성을 부여하고, 그 자체도 policy gradient loss로 업데이트되어 탐험 수준을 자동으로 조절한다. 논문이 말하는 “성공률이 100%로 수렴하면 노이즈가 자연히 감소”는, 사실상 학습 과정에서 exploration이 필요 없어질수록 \$sigma_{\theta'}$가 줄어드는 방향으로 최적화될 수 있다는 경험적 관찰이다.

 

그리고 재미있는 결정이 하나 더 있다. fine-tuning이 끝나면 noise net을 버리고 deterministic flow matching policy로 돌아간다고 한다. 이 말의 의미는 “초기 latent 샘플이 없다”가 아니라, 학습 중에 있던 per-step stochastic transition을 제거하고 ODE map을 되살린다는 의미로 이해하는 게 자연스럽다. 즉 학습은 SDE로, 실행은 ODE로. 여기서 잠재적 mismatch 의문(학습 때는 stochastic인데 실행 때는 deterministic)이 생기는데, 논문은 이를 일종의 training-time stochasticity(regularization)로 해석할 여지를 준다.

 

적당히 읽어보면 될 것 같다. 딱히 설명할 것이 없다.

 

알고리즘은 다음과 같다. 위 내용을 충분히 이해했으면 알고리즘 이해가 빠를 것이다:

Experiments

나머지 부분에 대해서는 내가 알고 싶었던 실험 결과가 아니고, 실험에서 보여주는 결과 자체도 논문의 아이디어와 크게 연관이 있는가 싶어서 다루지 않도록 하겠다...

Conclusion

ReinFlow는 솔직히 말해 “엄청 복잡한 새 알고리즘”이라기보다, flow 정책을 RL에 태우기 위해 필요한 최소한의 조건을 정확히 채운 설계로 보인다. 장점은 분명하고, 동시에 trade-off도 분명하다.

 

장점은 첫째로, flow 정책의 가장 큰 강점인 few-step/one-step 추론을 살리면서도, 그동안 발목 잡던 log-prob 계산 문제를 깔끔하게 제거했다는 점이다. CNF divergence 적분 + trace estimator는 이론적으로는 맞지만 로봇에서 요구하는 설정(큰 step size)에서 불안정해지기 쉬웠고, ReinFlow는 이를 “모델 정의 자체를 Gaussian Markov chain으로 바꾸는 방식”으로 해결한다. 이게 너무 정공법이라 오히려 강력하다. 둘째로, PPO 같은 안정적인 정책경사 프레임워크를 그대로 쓸 수 있게 만들어서 구현/연구 확장성이 높다. Markov process policy에 대한 policy gradient 정리를 끼워 넣은 덕분에, marginal을 직접 계산하지 않고도 joint log-prob를 사용해 업데이트할 수 있다. 셋째로, flow를 deterministic actor로 직접 미분할 때 자주 터지는 multi-step BPTT 체인 문제를 구조적으로 피한다(정확히는 pathwise derivative를 강제하는 형태를 피한다). RL fine-tuning이 “될 수 있는 형태”로 바뀌었다는 게 핵심이다.

 

단점/의문점은, 첫째로 이 접근이 “원래의 flow matching policy를 그대로 RL로 튜닝”한 것인지에 대해 엄밀히 말하면 애매하다. 학습 중 정책은 deterministic flow가 아니라 noise-injected Markov process다. 즉 RL이 보는 정책은 바뀌었다. 실행 시 noise net을 버리고 deterministic으로 복원한다고 해도, 학습 중 최적화된 목적이 execution-time 정책과 완전히 일치한다고 말하긴 어렵다. (물론 실전에서는 이런 train-test mismatch가 정규화로 작동하기도 하지만, 그건 장점이 아니라 성질이다.) 둘째로, noise injection이 “탐험을 원리적으로 설계한 메커니즘”인지에 대해서는 의심이 남는다. 이것은 ε-greedy 같은 명시적 탐험 정책과 다르고, reward가 sparse하거나 특정 방향으로 멀리 나가야 하는 경우에 local Gaussian perturbation이 충분한 탐험을 제공할지 불분명하다. 결국 탐험은 “support 주변 확산”에 가깝고, 얼마나 멀리/적절히 나갈지는 noise bound, entropy 정규화, 그리고 환경 특성에 많이 의존할 가능성이 크다. 셋째로, off-policy(SAC) 언급은 조심해서 봐야 한다. SAC를 “그대로” reparameterization pathwise gradient로 붙이면 multi-step chain에 대해 \partial a_K/\partial \theta를 다시 타게 되어 본래 피하려던 문제가 재등장할 수 있다. 논문이 제안하는 건 더 정확히는 Q를 weight로 쓰는 형태의 목적을 구성하거나(score-function 성격의 업데이트), off-policy로 쓸 수 있는 방향을 열어둔 것에 가깝고, naive하게 SAC actor gradient를 flow trajectory 전체로 backprop하는 건 위험하다.

 

결국 ReinFlow를 한 문장으로 요약하면 이렇다. flow를 RL에 올리기 위해 가장 치명적인 병목(확률 계산 불가/불안정)을 “학습 가능한 Gaussian 전이”로 해결하고, policy gradient가 작동하는 수학적 틀을 만들었다. 이게 장점의 전부이자 핵심이다. 동시에 “정책을 바꿔서 해결했다”는 점에서, 탐험의 질과 train-time stochasticity 제거 후의 성능 유지라는 질문이 자연스럽게 따라온다. 그래서 이 논문은 “flow를 RL로 튜닝할 수 있느냐”라는 질문에 대한 첫 번째, 꽤 설득력 있는 엔지니어링 답변으로 읽는 게 맞다고 본다.