Diffusion Policies Creating a Trust Region for Offline Reinforcement Learning 논문 리뷰

Summary

Diffusion Trusted Q-Learning (DTQL)은 diffusion 기반 behavior cloning policy가 정의하는 행동 매니폴드 위에서의 reward maximization 문제를 다루는 offline reinforcement learning 방법이다. 기존 diffusion 기반 offline RL 방법들은 diffusion 모델 자체를 정책으로 사용함으로써 높은 표현력을 얻는 대신, 반복적인 denoising sampling으로 인해 학습과 추론 비용이 크다는 한계를 가진다.

 

DTQL은 이 문제를 dual-policy 구조로 해결한다. 먼저 diffusion policy는 순수하게 behavior cloning만을 수행하여, 데이터 분포가 형성하는 복잡한 action manifold를 충실히 학습한다. 이후 실제로 학습·배포되는 target policy는 one-step Gaussian policy로 parameterize되며, diffusion policy로부터 유도된 diffusion trust region loss를 통해 데이터 매니폴드 내부에서만 탐색하도록 제약된다. 이 제약은 단순한 분포 정렬(KL)이 아니라, diffusion denoising loss에 기반한 기하학적 제약으로 작동하여 out-of-distribution 행동을 효과적으로 억제한다.

 

이러한 구조를 통해 DTQL은 diffusion 모델의 표현력을 유지하면서도, 빠른 추론 속도와 강한 mode-seeking 특성을 동시에 달성한다.

 

사실 Flow Q-learning(FQL)과 거의 유사하니 이 글을 읽는 독자들은 FQL과 이 논문과의 차이에 집중하기보다는, diffusion/flow 기반 offline RL에서 요구되는 것과 그것들이 아직 가지는 한계에 대해서 집중하며 보는 것이 좋을 것이라 본다.

Method

우선 policy 학습부터 확인해보자. 

여기서 Trust Region(TR) loss는 다음과 같다.

여기서 이미 학습된 $\mu_\phi$는 diffusion 기반 모델이며, 학습되는 policy는 $\theta$로 paramterize 되어있다.학습하고자 하는 policy가 action $a_\theta$를 출력했을 때, 해당 행동에 노이즈를 주입해서, diffusion BC policy가 복구하여 나오는 action과의 MSE loss를 줄이는 방향으로 학습된다. 즉 학습의 대상인 policy가 output하는 행동이, 실제로 데이터 내부에 존재하는 action인지에 대한 판단을 하여 기존 데이터 분포에서 너무 벗어난 action을 출력하는 경우 penalty를 주는 방식이다.

(사실 trust region이라고 그냥 이름을 지었을 뿐이지 그냥 distillation loss와 같다)

 

여기서 잘 알아두면 좋은 것은, 왜 KL divergence와 같은 제약식이 아닌 L2 norm으로 제약식을 두었는가에 관한 것이다. 이에 대한 실험 결과를 본 논문에서 보여주었는데, 확인해보자.

RL의 특수한 상황과 flow/diffusion으로 표현 가능한 support의 특수한 상황이 합쳐져서이다. KL divergence는 mode covering을 하도록 학습된다. 위의 figure을 보면 알 수 있듯이, 실제 데이터를 최대한 커버하려는 성향이 존재한다. 그러나 RL에서는 실제 행동의 선택에 대해서는 multimodality를 커버할 필요가 없는데, 왜냐하면 policy의 선택은 explicit하든 implicit하든 관계 없이 $\arg\max_a Q(s,a)$를 하기 때문이다.

 

즉 보통의 상황에서는 가장 큰 value를 가진 action만 취하면 되는데, 그렇다면 mode seeking만 하면 되는 것이지, 모든 mode를 충실히 커버할 필요가 없다는 것이다. 그림에서 보면 알 수 있듯이, KL은 실제로 존재하지 않는 데이터 영역에 대해서 커버를 하려고 한다. Offline Learning에서 data에 없는 state, action에 대한 value를 interpolation하게 되어 online 상황에서 학습되지 않은 action을 취하게 되는 가능성이 있어 문제가 된다는 것을 생각하면 보수적인 선택지를 가져가는 MSE loss가 더욱 합리적인 선택이 되는 것이다.

 

Q learning 자체는 우리가 알고 있는 IQL이다.

Experiments

 정도만 하면 충분할 것 같다. 사실 inference time이 짧은 것도... one-step policy라 당연한 것 아닌가 싶기는 하다 ㅋㅋ.

 

Limitations

사실 diffusion 기반이든 flow 기반이든 다를 것이 거의 없고, 그래서 본질적인 한계에 대해서 이야기해보고자 한다.

우선 이 논문이나 FQL이나 bc policy를 flow/diffusion으로 학습한 후, 실제 학습되는 target policy는 gaussian 또는 MLP로 parameterized 것을 공통적인 부분이다. 또한 flow/diffusion의 느린 generation으로 인해 target policy에 대한 아키텍처를 flow 또는 diffusion와 동일하게 할 수 없는 것도 사실이다.

 

이러한 점을 가졌을 때 1차적인 한계는 MSE loss가 optimal하냐는 것이다. 단적인 예시로, state에 대해서 offline data는 action을 정확히 두 개를 비슷한 양으로 가지고 있고, 그에 대한 보상이 같다고 해보자. 그런 경우 target policy는 그 중간, 즉 실제로 가지고 있지 않은 행동 영역으로 수렴해버린다. MSE loss의 자체적인 한계가 이 분야를 발전시킬 수 있는 하나의 연구 포인트인 것 같다.

 

근본적인 한계는 behavior policy는 diffusion/flow로 멋지게 만들어놓고, 한다는 것이 gaussian 또는 MLP policy를 타깃으로 삼는다는 것이다. 무언가 latency도 낮으면서 BC policy의 표현력 또는 manifold 자체를 더욱 잘 가져갈 수 있는 방법이 있지 않을까 하는 생각이 있다.

 

이만 논문 리뷰를 마치도록 하겠다.