[RL][Review] Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction (BEAR)

(논문의 의도를 가져오되, 개인적인 의견이 담길 수도 있습니다.)

Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction - Kumar et al, NeurIPS 2019 (논문, 코드)

요약

Off-policy RL은 샘플링 관점에서 효율적인 학습을 위해서 다른 policy (behavior policy)로부터 수집한 데이터로부터 경험을 배우는데 초점을 맞추지만, Q-learning이나 Actor-Critic 기반의 off-policy Approximate dynamic programming 기법은 학습시 사용된 데이터와 실제 데이터간의 분포가 다른 문제로 인해서 on-policy data를 추가로 활용하지 않고서는 성능을 개선하기가 어렵다.이 불안정성이 생기는 원인으로 Bootstrapping error를 꼽을 수 있는데, 이는 학습 데이터 이외에 놓여져 있는 action을 bootstrapping해올때 발생하는데, bellman backup operator에 해당 error가 계속 누적된다. 이를 극복할 수 있는 알고리즘으로 Bootstrapping Error Accumulation Reduction (BEAR)을 제안하는데, 이 알고리즘을 통해서 앞에서 언급한 off-policy distribution이 다른 상황에서도 robust하게 학습할 수 있다.

서문

대부분의 강화학습 알고리즘이 많은 데이터를 요구하기 때문에 데이터를 효율적으로 활용할 수만 있다면 실제 강화학습 문제에도 실용적으로 접근할 수 있고, 이전의 다양한 경험과 결합해서 조금 더 generalization할 수 있다. 일반적으로 off-policy RL이 이렇게 데이터를 활용할 수 있지만, off-policy data만 가지고 학습하기에는 제한되어 있다. off-policy data가 아무리 잘 되어 있다고 하더라도(해당 데이터를 전문가로부터 뽑았을지라도...) 문제는 존재하며, 여기에는 exploration에 대한 고려가 필요하다는 것을 나타낸다. 그래서 이 논문에서는 데이터가 많으면서 정적인 환경에서 off-policy value-based RL 알고리즘을 개발하고자 하는 것이다. 사실 off-policy 환경에 value-based 기법을 적용할 때 가장 큰 문제는 학습시 off-policy data를 활용하면서 Q-function이 out-of-distribution action, 즉 학습 분포에 없던 data를 가지고 평가가 되는 bootstrapping process에서 나타난다는 것이다. 이로 인해서 알고리즘이 새 데이터를 취합하지 못하면서, 학습이 불안정적이고, 발산할 가능성이 생긴다. 논문에서는 이 bootstrapping process에서 발생하는 error가 누적되는 현상에 대해 분석하고, 이 error를 처리할 수 있는 방법에 대해서 고민했다. 우선 off-policy data를 학습할 때 불안정적이고 성능이 낮게 나오는 원인에 대해서 공식화하고, 여기에서 action을 신중하게 선택한다면 앞에서 언급한 Q-function을 통해서 error가 전파되는 현상을 해소할 수 있다. 이렇게 bootstrapping error를 다루는 알고리즘을 Bootstrapping Error Accumulation Reduction (BEAR)라고 부르는데, error가 누적되는 것을 막기 위해서 support-set matching이라는 개념을 사용했다.

Out-of-Distribution Actions in Q-Learning

Q-learning 계열은 정적이면서 off-policy data 에서는 학습이 잘 되지 않는다. 우선 overfitting이 발생하는데, 그렇다고 static dataset의 크기를 늘린다고 해서 이 문제가 해결되지는 않는다. 이런 불안정성의 원인을 Bellman backup을 형성하는 부분에서 찾아볼 수 있다. Bellman error에 대한 제곱평균을 최소화하는 문제는 일종의 supervised regression 문제로 볼 수 있지만, 여기에서 regression의 target은 현재의 Q-function에 대한 추정치에서 도출된다. 이때 target은 next state에서 취할 action에 대한 학습된 Q value가 될텐데, 문제는 이 Q-function이 학습 데이터가 가지는 분포에서 나오는 입력에만 의존된다는 것이다. 결과적으로 단순히 Q-value를 최대화하는 과정은 \(\hat{Q}\) estimator가 학습 분포와는 거리가 있는 action들에 대해서 평가하게 되고 결과적으로 큰 error를 야기하는 요인이 된다. 여기서 이 action들을 Out-of-Distribution(OOD) action들이라고 표현한다.

이를 공식화해보면 \(\zeta_k(s, a) = \vert Q_k(s, a) - Q^*(s, a) \vert \) 를 Q-learning에서 \(k\)번 반복했을때의 total error라고 하고, \(\delta_k(s, a) = \vert Q_k(s, a) - \mathcal{T} Q_{k-1}(s, a) \vert \) 을 현재의 bellman error라고 해보자. (참고로 \(\mathcal{T}\)은 bellman operator) 그러면 이런 관계식을 가지게 된다.

$$ \zeta_k(s, a) \leq \delta_k(s, a) + \gamma \max_{a'} \mathbb{E}_{s'}[\zeta_{k-1}(s', a')] $$

그러면 \((s', a')\)에서의 error가 없어져도 \(k\)번째 반복에서 새로운 error인 \(\delta(s, a)\)가 계속 누적된다. OOD 에서의 state와 action에서의 error는 학습 데이터에서는 절대로 줄일 수 없기 때문에, \(\delta(s, a)\)는 클 것이다.

bootstrapping error를 줄이기 위해서 policy가 학습 데이터 분포의 support안에 놓여지는 action만 출력할 수 있도록 제한을 둘 수 있다. 이 방법은 학습시킬 policy의 distribution을 behavior policy와 유사하게끔 제한을 두는 BCQ 방식과는 차별되는 부분이다. 물론 BCQ에서 사용했던 방식은 action이 높은 확률로 학습 데이터 분포 내에 들게끔 할 수 있지만, 지나치게 제한적인 점도 있다. 예를 들어서 behavior policy가 uniform 분포를 가진다면, 학습된 policy는 거의 random하게 동작하게 되고, 아무리 좋은 policy를 만들기에 충분한 데이터를 가진다고 해도 좋은 성능을 내기 어렵다. 다시 말해서 behavior policy \(\beta(a \vert s)\)의 density가  일정 threshold 이상인 부분에서만 학습된 policy \(\pi(a \vert s)\)의 density 가 양의 값을 가지게 된다. 이를 공식으로 표현하면 다음과 같다.

$$ \forall a, \beta(a \vert s) \leq \epsilon \Rightarrow \pi(a \vert s) = 0 $$

여기에서 제한이 너무 과도할 경우 data distribution에 머무는 데이터를 뽑는 것과 suboptimal solution을 찾는 것 사이에 tradeoff가 존재한다는 것을 찾았고, 이를 통해서 학습된 policy의 support에 제한을 두되, 이미 support안에 포함된 action에 대한 확률에는 두지 않는 방법을 생각했다. 이 방법을 통해서 OOD action에서 Q-function estimator \((\hat{Q})\)를 평가하는 것은 피하면서, 효율적인 정책을 찾을 수 있는 유연함은 유지할 수 있다.

Distribution-Constrained Backups

앞에서 언급한것처럼 policy가 학습 데이터의 분포내의 action만 출력할 수 있도록 제한을 두는 과정이 필요한데, 논문에서는 이에 필요한 과정을 Distribution Constrained Backup이라고 표현했다. 만약 어떤 policy들의 집합 \(\Pi\)가 있을때, Distribution-Constrained Backup은 다음과 같이 정의할 수 있다.

$$ T^{\Pi}Q(s, a) \stackrel{\text{def}}{=} \mathbb{E}\Big[R(s, a) +\gamma \max_{\pi \in \Pi} \mathbb{E}_{P(s' \vert s, a)}[V(s')]\Big] \qquad V(s) \stackrel{\text{def}}{=} \max_{\pi \in \Pi} \mathbb{E}_{\pi}[Q(s, a)]$$

여기에서 backup operator \(T^{\Pi}\)는 우리가 알고있는 일반적인 bellman backup의 특성을 그대로 가지고 있다. 만약 어떤 값을 근사했을 때의 error에 대한 이 backup의 operator의 optimality를 파악하기 위해서 error가 발생하는 두가지 요소에 대해서 정량화를 했다. 하나는 suboptimality bias이고, 다른 하나는 학습 데이터 분포와 backup에 사용되는 policy들의 분포간의 차이에서 발생하는 것인데, 이 두개를 통해서 앞에서 언급한 OOD action들을 표현할 수 있다. 여기서 우리가 최종적으로 구한 solution의 suboptimality를 파악하기 위해서 suboptimality constant라는 지표를 정의해서 optimal policy \(\pi^*\)가 \(\Pi\)에서 얼마나 멀어져 있는지를 정의하고자 했다.

$$ \alpha(\Pi) = \max_{s, a} \vert T^{\Pi}Q^*(S, a) - T Q^*(s, a) \vert $$

그리고 concentrability coefficient를 정의해서 \(\Pi\)로부터 생성된 visitation distribution이 학습 데이터 분포와는 얼마나 떨어져 있는지를 측정했다. 결국 어떤 state와 action이 distribution으로부터 떨어져 있는 정도를 나타낸 것이다.

Bootstrapping Error Accumulation Reduction (BEAR)

결국 앞에서 언급한 Distribution Constrained Backup을 기존의 TD3나 SAC와 같은 Actor-Critic 알고리즘에 적용해서 bootstrapping error가 누적되는 것을 줄이고자 했다. 여기서 핵심은 학습 데이터 분포와 같은 support를 가지는 policy를 찾으면서, 갑작스럽게 error가 누적되는 현상을 막는 것이다. 알고리즘은 크게 두 부분으로 구성되어 있는데, BCQ처럼 \(K\)개의 Q-function을 가지고 이 중 가장 작은 Q-value를 policy improvement하는데 사용했고, 또한 behavior policy와 같은 support를 공유하는 policy를 policy 의 집합 \(\Pi_{\epsilon}\)내에서 찾을 수 있는 constraint를 설계했다. 이 과정들 모두 기존의 actor-critic 알고리즘 내의 policy improvement 단계를 수정한 것이다. 그리고 \(K\)개의 Q-function에 대한 평균값으로도 policy improvement가 수행될 수 있고, 이 구조도 좋은 성능을 보여준다.

우선 \(K\)개의 Q-function을 다음과 같이 정의한다: \(\hat{Q}_1, \ldots, \hat{Q}_K\). 그리고 \(\Pi_{\epsilon}\)내에서 Q-value에 대한 conservative estimate를 최대화할 수 있도록 policy update한다.

$$ \pi_{\phi}(s) \mathrel{:=} \max_{\pi \in \Pi_{\epsilon}} \mathbb{E}_{a \sim \pi(\cdot \vert s)} [ \min_{j=1,\ldots, K} \hat{Q}_j(s, a)] $$

실제 적용할때는 behaviour policy \(\beta\)는 모르는 상황이기 때문에, \(\pi\)를 \(\Pi\)내에 제한을 두기 위한 근사방법이 필요하다.이를 위해서 미분가능한 제한값을 둬서 \(\pi\)가 \(\Pi\)내에 대략적으로 제한을 둔후 dual gradient descent를 통해서 제한된 optimization problem을 대략적으로 풀었다. 참고로 \(\beta\)와 actor \(\pi\) 사이의 maximum mean discrepancy (MMD)의 샘플링된 버전을 사용했다. (MMD란 주어진 두개의 분포사이의 거리를 통해서 분포가 똑같은지 여부를 판단할 수 있는 kernel 기반의 통계기법을 말하는데, density estimation을 할 필요가 있는 상황에서는 loss function으로 활용할 수도 있다.)

만약 \(x_1, \dots, x_n \sim P\)이고 \(y_1, \dots, y_m \sim Q\)라고 샘플들이 주어져 있을때, \(P\)와 \(Q\)사이의 MMD는 다음과 같이 계산할 수 있다.

$$ \text{MMD}^2( \{x_1, \dots, x_n\}, \{y_1, \dots, y_m\}) = \frac{1}{n^2}\sum_{i, i'}k(x_i, x_{i'}) - \frac{2}{nm} \sum_{i, j}k(x_i, y_j) + \frac{1}{m^2}\sum_{j, j'}k(y_j, y_j') $$

여기서 \(k(\cdot, \cdot)\)은 어떠한 kernel이라도 상관이 없는데, 실험상에서는 Laplacian kernel과 Gaussian kernel이 잘 동작했다. MMD 공식은 두 분포의 density와는 크게 연관이 없어서, 샘플을 통해서 직접적으로 optimize할 수 있다. 경험적으로 찾은 결론으로는 \(P\)와 \(Q\)사이의 샘플링된 데이터의 MMD는 사실 \(P\)의 support를 가지는 uniform distribution과 \(Q\)사이의 MMD와 유사하다는 점인데, 이 말은 결국 주어진 support상에서 분포를 제한할때 MMD를 활용할 수 있다는 것이다.

앞에서 말한 과정을 모두 적용해보면, policy improvement step에서의 optimization problem은 다음과 같이 정의할 수 있다.

$$ \pi_{\phi} \mathrel{:=} \max_{\pi \in \Delta_{\vert S \vert}} \mathbb{E}_{s \sim \mathcal{D}} \mathbb{E}_{a \sim \pi(\cdot \vert s)} \Big[ \min_{j=1, \ldots, K} \hat{Q}_j(s, a)\Big] \quad \text{ s.t. } \mathbb{E}_{s \sim \mathcal{D}} [ MMD(\mathcal{D}(s), \pi(\cdot \vert s))] \leq \epsilon $$

\(\epsilon\)은 대략적으로 정한 threshold이다.

Bootstrapping Error Accumulation Reduction (BEAR)

그러면 BEAR가 앞에서 언급한 distribution-constrained backup과는 어떻게 연결되어 있을까? 위의 

...

개인생각

Levine교수 랩에서 나온 offline RL논문인데, 일단 논문이 조금 어렵다. 우선 distribution에 constraint를 주기 위한 수학적 증명이 쭉 되어 있는데, 이를 이해하기 위한 기반 지식이 조금 필요할 듯 하다. 무엇보다 왜 constraint를 줘야하는지에 대한 설명이 뒤의 appendix에까지 이어져 있어서, 누구한테 이 부분을 설명하기에는 먼저 필요 지식부터 이해하는게 중요할 듯 하다. (특히 MMD는 kernel에 대해서 이해를 먼저 해야할 것 같다.) 다행인 점은 저자가 블로그를 통해서 조금 쉽게나마 설명해줬다는 점? 사실 코드를 봐도 논문에 언급된 알고리즘이 명확하게 보이지 않아서 이해가 좀 어렵다.