Reversal Q-Learning: Flow-Matching 정책과 함께 작동하도록 Offline RL을 학습시키는 방법

Flow matching은 생성 모델링 (generative modeling) 도구 상자에서 가장 유용한 도구 중 하나가 되었습니다. 이는 확산 모델 (diffusion models)보다 빠르게 학습되고, 고품질의 샘플을 생성하며, 다중 모드 분포 (multimodal distributions)를 잘 처리합니다. 이러한 특성은 주어진 상황에서 "올바른" 움직임이 단일 지점이 아니라 일련의 그럴듯한 행동 집합일 수 있는 로봇 행동 모델링에 매우 매력적입니다.

문제는 flow matching을 강화학습 (reinforcement learning)과 결합하는 것이 진정으로 어렵다는 점이며, 특히 새로운 경험을 수집할 수 없고 고정된 데이터셋만 있는 offline setting (오프라인 설정)에서는 더욱 그렇습니다. Aditya Oberai, Seohong Park, Sergey Levine의 새로운 논문인 Reversal Q-Learning (RQL)은 이 문제에 대한 깔끔한 해결책을 제안하며, 그 핵심 아이디어는 자세히 이해할 가치가 있을 만큼 우아합니다.

왜 Flow Matching과 Offline RL은 서로 잘 맞지 않는가

RQL이 해결하는 문제를 이해하려면 flow matching이 실제로 무엇을 하는지 아는 것이 도움이 됩니다. Flow matching 정책은 단순한 노이즈 분포 (noise distribution)에서 타겟 행동 분포 (target action distribution)로 샘플을 운송하는 벡터장 (vector field)을 학습합니다. 추론 (inference) 시점에는 노이즈에서 시작하여 F 단계 동안 벡터장을 적분하여 행동을 생성합니다. 단계가 많을수록 정책의 표현력 (expressive)은 높아지지만, 계산량도 늘어납니다.

강화학습을 사용하여 이 정책을 개선하고자 할 때, 하위 리턴 (downstream returns)을 기반으로 행동에 대한 크레딧 (credit)을 할당해야 합니다. Offline RL에서는 정적 데이터셋에서 추정된 Q-함수 (Q-functions)를 사용하여 이를 수행합니다. 문제는 데이터셋이 가공되지 않은 (state, action) 쌍을 포함하고 있다는 점입니다. 즉, 해당 행동을 생성한 중간 flow 단계 (flow steps)에 대한 기록이 없습니다. flow 단계는 보이지 않습니다.

이를 해결하기 위한 하나의 원칙적인 방법은 확장된 MDP (expanded MDP) 프레임워크를 사용하는 것입니다. 즉, F개의 flow 정제 단계 (flow refinement steps) 각각을 더 긴 마르코프 결정 과정 (Markov decision process)에서의 별도 행동 (action)으로 취급하는 것입니다. 이렇게 하면 flow 단계가 명시적으로 드러나며 표준적인 Q-learning을 적용할 수 있습니다. 하지만 이는 두 가지 새로운 문제를 야기합니다.

1. 데이터셋 불일치 (Dataset incompatibility). 여러분의 오프라인 데이터셋에는 중간 flow 상태 (intermediate flow states)가 포함되어 있지 않습니다. 데이터에 존재하지 않는 전이 (transitions)에 Q-learning을 직접 적용할 수는 없습니다.

2. 호라이즌의 저주 (The curse of horizon). MDP를 F배로 확장한다는 것은 유효한 계획 호라이즌 (planning horizon)이 F배로 늘어남을 의미합니다. 시간차 (Temporal difference, TD) 학습은 긴 호라이즌에 대해 편향 (bias)을 누적시키므로, 가치 추정치 (value estimates)의 신뢰도가 떨어지게 됩니다.

이전의 접근 방식들은 가중 회귀 (weighted regression), 증류 (distillation), 또는 거부 샘플링 (rejection sampling)을 사용하여 이러한 문제들을 우회해 왔으나, 이 방법들은 모두 정보를 버리거나 자체적인 근사 오차 (approximation errors)를 유발한다는 단점이 있습니다.

RQL의 해결책: 발생한 일을 재구성하기 위해 Flow를 역전시키기

RQL의 핵심 통찰은 결정론적 flow ODE (deterministic flow ODEs)는 **가역적 (reversible)**이라는 점입니다. 만약 상태 s에 대해 정책이 생성한 최종 행동 a를 알고 있다면, flow ODE를 역방향으로 실행하여 그 행동에 이르게 된 중간 상태들의 전체 시퀀스 x⁰, x¹, ..., xᶠ를 복구할 수 있습니다.

공식적으로, 오프라인 데이터셋의 임의의 전이 (s, a, r, s')에 대해, RQL은 다음의 역방향 ODE를 해결합니다:

d/df θ(s, x, f) = -v(s, θ(s, x, f), f)

여기서 v는 학습된 벡터장 (vector field)입니다. 이는 flow 공간을 통과하는 "가상"의 온-폴리시 궤적 (virtual on-policy trajectory)을 재구성합니다. 즉, 현재 정책이 상태 s로부터 행동 a를 생성하기 위해 거쳤을 중간 상태들의 정확한 시퀀스입니다.

이 가상 궤적들은 현재 flow 정책에 대해 **결정론적 (deterministic)**이며 **온-폴리시 (on-policy)**적입니다. 이것이 바로 이 궤적들이 유용한 이유입니다. 온-폴리시적이기 때문에, 오프-폴리시 편향 (off-policy bias)을 유발하지 않으면서 flow 단계 전체에 걸쳐 다단계 리턴 (multi-step returns)을 적용할 수 있습니다. 또한 결정론적이기 때문에, 다단계 리턴은 샘플링된 추정치가 아닌 정확한 값(exact)이 됩니다.

호라이즌 축소 (Collapsing the Horizon)

두 번째 혁신은 호라이즌의 저주 (curse of horizon) 문제를 다룹니다. 가상 궤적 (virtual trajectories)은 결정론적 (deterministic)이며 온-폴리시 (on-policy) 방식이기 때문에, RQL은 다단계 리턴 (multi-step returns)을 사용하여 중간의 flow 단계들을 완전히 건너뛸 수 있습니다. T × F 단계의 호라이즌 (T는 태스크 호라이즌, F는 flow 단계의 수)에 대해 가치 함수 (value function)를 추정하는 대신, RQL은 유효 호라이즌을 다시 T로 축소 (collapse)합니다.

이는 중간 flow 단계들이 환경과 상호작용하지 않기 때문에 가능합니다. 이 단계들은 순수하게 정책의 생성 과정 내부에만 존재합니다. 보상 신호 (reward signal)는 각 flow 단계 이후가 아니라, 전체 액션이 완료된 시점에만 전달됩니다. 따라서 가치 추정 (value estimation) 목적상 전체 flow 생성을 하나의

구현은 JAX로 되어 있으며 GitHub에서 확인 가능하므로, Offline RL (오프라인 강화학습) 분야에서 연구하는 연구자들이 비교적 쉽게 접근할 수 있습니다.

로보틱스에서 이것이 중요한 이유

Offline RL (오프라인 강화학습)은 로보틱스 분야에서 특히 중요한데, 그 이유는 온라인 경험 (online experience)을 수집하는 것이 비용이 많이 들고, 느리며, 때로는 안전하지 않기 때문입니다. 로봇 시연 (robot demonstrations)으로 구성된 대규모 데이터셋은 — 설령 그것이 불완전할지라도 — Offline RL이 단순히 행동을 모방하는 것이 아니라 보상 (reward)을 최적화함으로써 시연보다 개선된 정책 (policy)을 추출할 수 있게 해줍니다.

Flow matching (플로우 매칭)은 로봇 정책에 있어 매력적인데, 로봇의 행동은 종종 다봉성 (multimodal)을 띠기 때문입니다. 예를 들어, 물체를 잡는 데에는 여러 가지 동등하게 유효한 방법이 있을 수 있으며, 단봉형 가우시안 정책 (unimodal Gaussian policy)은 이들을 평균 내어 유효하지 않은 행동을 만들어낼 수 있습니다. RQL은 이전 방법들이 요구했던 근사 (approximations) 없이도 표현력이 풍부한 flow 정책을 Offline RL과 결합하는 것을 실용적으로 만들어 줍니다.

더 넓은 맥락

RQL은 RL을 이용한 생성 정책 학습 (training generative policies with RL)에 관한 점증하는 연구 흐름에 부합합니다. 관련 접근 방식으로는 GenPO (on-policy RL을 위해 정확한 diffusion inversion을 사용)와 FMER (flow 정책과 함께 advantage-weighted regression을 사용) 등이 있습니다. RQL을 차별화하는 점은 Offline RL (오프라인 설정)에 초점을 맞추고 있으며, ODE 가역성 (ODE reversibility)을 사용하여 데이터셋 불일치 문제를 완전히 회피한다는 것입니다. 확장된 MDP 프레임워크 자체는 새로운 것이 아니지만, 이를 오프라인에서 적용하기 위해서는 RQL이 도입한 가상 궤적 구축 (virtual trajectory construction)이 필요했습니다.

요약

Reversal Q-Learning은 구체적인 기술적 장애물, 즉 중간 flow 상태를 포함하지 않는 오프라인 데이터를 사용하여 flow-matching 정책에 Q-learning을 어떻게 적용할 것인가라는 문제를 해결합니다. 그 해결책은 — flow ODE를 역방향으로 실행하여 가상 on-policy 궤적을 재구성한 다음, multi-step returns를 사용하여 확장된 horizon을 축소하는 것 — 은 기술적으로 깔끔하며 경험적으로 효과적입니다. 생성 모델 (generative models)과 Offline RL (오프라인 강화학습)의 교차점에서 연구하는 연구자들에게 이는 유용한 도구가 될 것입니다.

이 논문은 arxiv.org/abs/2606.17551에서 확인할 수 있으며, 코드는 github.com/aoberai/rql에서 제공됩니다.