ReFORM: Reflected Flows for On-support Offline RL via Noise Manipulation 논문 리뷰

Summary

본 논문은 Offline Reinforcement Learning에서 활용되는 KL 기반 trust-region 기법과 Wasserstein regularization이 분포 간 거리를 제어함에도 불구하고, 정책의 support 자체를 엄밀히 제한하지 못한다는 점에 주목한다. 이러한 구조적 한계는 데이터 분포 외 영역으로의 정책 질량 이동을 완전히 방지하지 못하며, 그 결과 OOD 행동에 대한 가치 추정 오류가 누적되는 문제가 발생한다. 이에 본 연구는 정책의 support를 직접적으로 제약하는 구조를 설계하고, 이를 통해 OOD action에 대한 이론적·실증적 제어 가능성을 분석한다.

 

Preliminaries

Offline Reinforcement Learning

Offline RL은 online RL과 다르게 고정된 데이터셋 $D=\{(s_i, a_i, r_i, s'_i)\}_{i=1}^N$만을 이용하여 정책을 학습하는 문제로 정의된다. 따라서 환경과의 추가적인 상호작용 없이 정책을 최적화해야 하므로, 데이터 분포 밖의 행동(out-of-distribution actions)에 대한 가치 함수의 외삽 오류(extrapolation error)가 주요한 실패 원인으로 알려져 있다.

 

이를 완화하기 위해 많은 기존 연구들은 정책이 데이터 행동 분포에서 크게 벗어나지 않도록 분포 제약(distributional constraint)을 도입한다. 이러한 방법으로 잘 알려진 KL과 Wasserstein 기반 regularization이 있는데, 간단하게 살펴보자:

 

1. KL 기반 Regularization

정책의 업데이트는 다음과 같은 제약식을 바탕으로 진행된다:

$$\max_{\pi} \ \mathbb{E}_{s \sim \mathcal{D},\, a \sim \pi}[Q(s,a)] \quad \text{s.t.} \quad \mathrm{KL}(\pi(\cdot|s)\,\|\,\pi_{\text{BC}}(\cdot|s)) \le \epsilon$$

이러한 접근은 정책을 behavior policy 근처에 머물도록 함으로써 OOD 행동 생성을 억제하는 것을 목표로 둔다. 그러나 본 논문에서는 다음과 같은 정리로 KL divergence는 support 제약을 과하게 적용함을 보인다:

즉 KL 제약은 target policy의 support가 behavior policy의 support 안에 들도록 하는 것을 유도하지만, support 제약을 만족함에도 KL 값이 unbounded일 수도 있는 것이다. Support 제약이 근본적인 목표이지만 사실 support 제약을 넘어 density ratio 제약도 동시에 거는 필요 이상의 제약이며, 또한 실제 연속 제어 설정에서는 support 제약을 실질적으로 제공하지도 못한다.

 

2. Wasserstein 기반 Regularization

정책 업데이트는 다음과 같다:

$$\begin{aligned} \max_{\pi} \quad & \mathbb{E}_{s \sim \mathcal{D}, a \sim \pi(a|s)} [Q(s, a)] \\ \text{s.t.} \quad & \left( \inf_{\gamma \in \Pi(\pi, \pi_{\text{BC}})} \int_{\mathcal{A} \times \mathcal{A}} \|a - a'\|^p \, d\gamma(a, a') \right)^{1/p} \le \epsilon \end{aligned}$$

 

Wasserstein distance 기반 제약은 분포 간 질량 이동 비용을 직접 제어한다는 면에서 KL보다 기하학적 의미가 명확한 것은 사실이나, 본 논문에서는 다음과 같은 proposition을 제공하여 Wasserstein도 본질적인 면에서 올바르지 않음을 보인다:

즉, 작은 Wasserstein 거리 하에서도 정책은 데이터가 존재하지 않는 영역에 소량의 질량을 배치할 수 있다. Offline RL에서는 이러한 소량의 질량이라도 Q-function의 외삽 오류가 클 경우 정책 최적화 과정에서 증폭될 수 있다.

 

Method

사실 Offline RL의 시초가 되는 논문에서도, "사실 support에 대한 제약을 두고 싶었는데 그러지 못해 다른 제약(그당시에는 MMD)을 걸었다"고 한 바가 있다. 본 논문에서는 이러한 제약을 어떻게 걸었는지에 대해서 알아보도록 하자:

그러나 이런 제약을 두는 것은 말로는 쉽지, 실제로 구현에 있어서는 너무 어렵다. 왜냐하면 우선 behavior policy의 support를 추정하는 것은 복잡하거나 어려우며, 그 support 안에 $\pi_\theta$의 action을 가두는 것도 쉬운 문제가 아니다. 즉 support를 학습하면서도 강제하는 문제가 동시에 생기는 것이다. 이 논문에서는 이러한 문제를 어떻게 해결하는지 알아보자.

 

ReFORM의 핵심 아이디어는 support를 직접 추정/정규화로 강제하지 않고, 정책 파라미터화 자체를 support-safe하게 만드는 것이다. 이를 위해 먼저 behavior action 분포를 모사하는 BC flow policy $\mu_{\theta_1}$를 학습한다. 중요한 점은 flow의 source distribution $q_{\text{BC}}$를 bounded support를 갖도록 선택하는 것이다. 예를 들어,

$$q_{\text{BC}} = U(B_l^d), \quad \mathrm{supp}(q_{\text{BC}}) = B_l^d.$$

BC flow가 $z \in B_l^d$를 action으로 매핑한다고 하면, $z$가 $B_l^d$ 안에 있는 한 생성된 행동은 자연스럽게 BC가 모델링한 support 영역 안에 놓이게 된다. 즉 “행동 공간에서 support를 직접 다루지 말고, bounded latent에서 최적화하자”가 1차 결론이다.

 

BC만으로는 데이터 평균 성능에 묶이므로, ReFORM은 “행동”을 직접 바꾸지 않고 latent noise 분포를 바꿔서 성능을 올린다. 구체적으로 $w \sim U(B_l^d)$에서 시작해 noise generator $\mu_{\theta_2}$가 $z = \mu_{\theta_2}(w;s)$를 생성하고, 최종 행동은

$$a = \mu_{\theta_1}(z;s)$$

로 만든다. 따라서 우리가 최적화하는 대상은 사실상 “행동”이 아니라 “$z$의 분포”다. 이때 반드시 필요한 조건은:

$$\mathrm{supp}(\tilde q_{BC}) \subseteq \mathrm{supp}(q_{BC}) = B_l^d$$

즉, noise manipulation 이후에도 latent가 ball 밖으로 나가면 안 된다.

 

간단하게 말하면, 원래는 latent noise에서 action을 생성하지만, 우리는 noise를 compact한 bounded source에서 샘플링하고, noise manipulation 이후에도 latent 분포의 support가 원래 source support를 벗어나지 않도록 강제한다. 그러면 BC flow 정책이 동일한 mapping으로 latent를 action으로 pushforward하므로, latent에서의 support 포함 제약이 action space에서도 support 포함 제약을 간접적으로 보장한다.

 

그러나 이러한 latent noise를 생성함에 있어서 일반적인 flow 기반 생성 모델은 연속적인 벡터장을 따라 ODE를 적분하여 분포를 변형하므로, 그 궤적이 경계를 벗어나는 것을 구조적으로 막지 못한다. 즉, 표현력은 충분하지만, support를 보존하지는 않는다. 따라서 논문에서는 reflected flow를 사용하여 표현력 유지와 domain 내부에 가두는 동역학을 동시에 달성한다:

Algorithm

그림과 함께 알고리즘을 이해하면 편하다. Critic update에서 one-step policy의 action sampling에서 noise는 단순 랜덤 노이즈가 아닌 noise flow 모델의 노이즈이며, 그러한 노이즈 생성 모델은 Q maximize를 위해 학습이 되는 것이 특징이다.

 

여기서 논문이 밝히는 limitation이 나오는데, noise flow 모델이 학습되는 과정에서 BPTT가 필요하다는 것이다.

 

Conclusion

Support constraint를 영리하게 제어했다는 점에서 이 논문이 Offline RL의 하나의 milestone이 아닐까 싶다. 다만 아직 support constraint가 latent에 대한 것이기 때문에, action에 대한 constraint가 아니기 때문에 불완전하다는 점, 그리고 support constraint은 잘 건드렸지만 offline RL에서는 support에 포함되어도 데이터 밀도가 낮거나 학습이 덜 된 영역에서는 아직도 위험하다는 점이 한계이다.

 

그리고 진짜 one-step은 아니다. BPTT로 인한 학습 저하도 하나의 문제지만, one-step policy에게 넘겨주는 latent 자체가 noise flow 모델에 의해서 제공되기 때문에 구조적으로 따로 이득이 있어보이지는 않는다.