[RL][Review] Offline RL without Off-Policy Evaluation (onestep-rl)

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

Offline RL without Off-Policy Evaluation - Brandfonbrener et al, NeurIPS 2021 (논문, 코드)

요약

이전에 수행된 대부분의 Offline RL에서는 off-policy evaluation과 관련된 반복적인 Actor-critic 기법을 활용했다. 이 논문에서는 behavior policy의 on-policy Q estimate를 사용해서 제한된/정규화된 policy improvement를 단순히 한번만 수행해도 잘 동작하는 것을 확인했다.이 one-step baseline이 이전에 발표되었던 논문에 비하면 눈에 띌만큼 간단하면서도 hyperparameter에 대해서 robust한 결과를 얻을 수 있었다. 이전의 iterative 방법들이 상대적으로 낮은 성능을 띄는 이유는 off-policy evaluation을 하면서 high variance가 상속되고, 이 estimate에 대해서 policy에 대한 최적화 작업이 반복되면서 점점 커지기 때문이었다.

서문

효율적인 realworld RL로 나아가는 중요한 단계 중 하나는 바로 sample efficiency를 향상시키는 것인데, 이 중 하나의 접근 방식이 Batch RL이라고도 알려져 있는 Offline RL인데, 이 방법은 환경에 대한 직접적인 interaction없이 다른 behavior policy에 의해서 수집된 data만으로 새로운 policy를 학습하는 것이다. 이전까지 발표된 deep offline RL 알고리즘들은 policy evaluation과 policy improvement를 번갈아 하는 actor-critic style을 띄고 있다. 그리고 이 모든 알고리즘들이 critic을 학습시킬때 off-policy evaluation에 매우 의존하게 된다. 그래서 논문에서는 behavior Q function을 사용해서 한번만 policy improvement를 하더라도 앞에서 언급한 iterative 알고리즘보다 성능이 좋다는 것을 보여주고자 한다.보통 iterative 알고리즘들이 성능이 안 좋은 이유는 off-policy evaluation이 안좋게 수행됨으로써, 부정확한 Q value를 만들기 때문인데, 이에 대한 원인으로는 다음과 같이 두가지를 들 수 있다.

1. behavior policy와 평가될 policy (target policy)간의 distribution shift
2. policy optimization으로 인해서 bias가 생기고, dynamic programming을 통해서 이 bias가 state space로 발산되는 iterative error exploitation

물론 one-step 알고리즘이 항상 좋은 성능을 보여주는 것은 아닌데, 뒷부분에서 iterative 알고리즘이 one-step 알고리즘보다 더 좋은 성능을 보여주는 경우에 대해서 몇가지 언급하고자 한다. 간단하게 말하면 dataset이 매우 커서 behavior policy가 state-action space 상에 잘 수렴하고 이로 인해서 off-policy evaluation 가 잘 수행될 때는 iterative 알고리즘이 더 효과적일 수 있다.

그림 1. one-step과 multi-step의 차이

참고로 위의 그림은 논문에서 지칭한 one-step과 multi-step의 차이를 도식화한 것이다. 가령 one-step을 취하는 경우 behavior policy \(\beta\)에서 다음 policy \(\pi_1\) 수행할 때는 on-policy \(\hat{Q}^\beta\)를 사용하는데 비해, multi-step은 \(\pi_1\)이후에도 해당 policy의 Q estimate, 즉 off-policy \(\hat{Q}^{\pi_i}\)를 사용하는 것이다. 이렇게 보면 one-step은 multi-step보다 상대적으로 behavior policy \(\beta\) 주변에 있는 "safe" policy를 찾게 되는, 일종의 제한이 걸린 셈이고, multi-step은 이제 behavior policy와는 거리가 있는 상대적으로 "unsafe" policy로 Q estimate가 계산될 것이다. 어떻게 보면 multi-step의 의미속에는 앞에서 언급한 iterative의 내용이 담겨 있는 것이다.

Iterative algorithms

앞에서 언급했다시피 대부분의 deep offline RL 알고리즘은 iterative actor-critic으로 이뤄져 있다. 해당 알고리즘들은 학습된 policy가 behavior policy에 의해서 생성된 data로부터 너무 멀리 떨어지지 않게끔 하기 위한 다양한 방법들을 제안했다. 크게 3개의 그룹을 두자면, policy constraints/regularization과, imitation learning에 대한 수정, 그리고 Q값에 대한 regularization 방법들이다.

1. policy constraints/regularization
   : 이 방법들은 policy에 직접적으로 제한을 두는 방식이다. 몇몇 저자들은 \((s, a)\)가 데이터를 생성하는 분포 상에 존재하도록 충분한 support를 가지게끔 action을 선택할 수 있게 학습된 policy에 대한 제한을 두었다. 또다른 방식은 학습된 policy를 behavior policy를 향하게끔 regularization을 취하는 것이다. 보통은 KL divergence나 MMD 방식을 취한다.  이 방식들은 데이터들이 behavior policy의 분포와 얼마나 근접한지를 나타내주는 hyperparameter를 가지는 매우 직관적인 방식이 되겠다. 그리고 iterative하고, off-policy evaluation에 의존한다.
2. modification of imitation learning
   : 어떤 알고리즘들은 낮은 Q value를 필터링하고, 나머지 높은 Q value를 가지고 imitation learning을 취하는 방식을 택했다. 또다른 방식으로는 Q value에 대해서 지수를 취한 값을 weight으로 가지는 weighted imitation learning 방식을 사용하기도 했다. 이 알고리즘들도 역시 iterative하다.
3. Q regularization
   : 학습된 policy가 미지의 action을 선택하는 것을 막을 수 있는 방법 중 하나는 behavior policy에 인접하도록 도와주는 regularization 기법을 사용하고, 미지의 state, action 쌍에 대해서는 비관적인 평가값을 가지게 하는 것이다. 하지만 미지의 state에 대해서 불확실성을 적절하게 정량화하는 것은 신경망의 value function을 다룰때 매우 어렵게 작용한다.

One-step algorithm

최근 발표된 논문에서는 behavior value function을 기반으로 policy를 최적화해도 꽤 잘 동작하는 모습을 보여주고 있다. 어떤 논문에서는 D4RL 벤치마크에서 continuous control task에 대해서 연구를 진행했는데, 이 논문에서는 ensemble, distributional Q function을 사용하고, 최신 regularization 기법을 활용했다. 또 다른 논문에서는 behavior value estimation이라고 하는, discrete action 환경에서 one-step 알고리즘을 수행하는 것에 대해 연구했는데, 성능이 Atari game이나 RL Unplugged 벤치마크에 있는 discrete task에서 좋게 나왔다. 또한 evaluation 단계에서 활용할 수 있는 최신 regularization에 대해서도 소개했다. 사실 discrete 환경과 continuous 환경에서 돌린 것과의 차이가 있는데, continuous 환경을 위해서는 parametric policy들로 구성된 actor-critic 알고리즘이 필요한 반면, discrete 환경에서는 policy improvement 단계에서 Q function을 정확하게 계산할 수 있다는 점이다. 이 논문에서는 iterative 알고리즘의 성능이 안좋은 원인을 overestimation으로 정의했는데, 우리는 이 overestimation을 유발하는 원인으로 distribution shift와 iterative error exploitation을 분리해서 정의했다. 이렇게 분리를 해서 보게 되면 iterative 알고리즘에 의해서 생기는 특정 문제 상에서의 off-policy evaluation의 근본적인 한계에 대해서 파악하는데 도움이 되고, 어떻게 보면 앞으로 진행될 일의 근거를 줄 수 있다.

Algorithm 

Algorithm template

우선 일반적인 offline approximate modified policy iteration (OAMPI) 라는 구조를 고려했다. 알고리즘은 아래와 같다.

그림 2. Offline Approximate Modified Policy Iteration (OAMPI)

알고 있다시피 Policy Iteration은 크게 두가지로 나눠진다. 바로 Policy evaluation과 Policy improvement. Policy evaluation에서는 dataset \(\mathcal{D}_N\)만 가지고 현재의 policy \(\pi_{k-1}\)에 대한 Q function의 estimate \(\hat{Q}^{\pi_{k-1}}\)을 구한다. approximate process를 warm-start하기 위해서 이전의 Q estimate \(\hat{Q}^{\pi_{k-2}}\)도 사용하기도 한다. 두번째 단계에서는 evaluation 단계에서 구한 \(\hat{Q}^{\pi_{k-1}}\)와 estimated behavior policy \(\hat{\beta}\), 그리고 dataset \(\mathcal{D}_N\)을 사용해서 새로운 policy \(\pi_k\)를 찾게 되는 policy improvement 과정을 수행한다. 이 과정에서도 역시 warm-start를 위해서 이전의 policy \(\pi_{k-1}\)을 사용한다. 여기에다가 \(\pi_k\)가 \(\beta\)와 \(\mathcal{D}_N\)의 support 상에 존재하도록 regularized하거나 constraint를 부여하는 과정이 포함될 수 있다.

One-step

가장 간단한 알고리즘은 바로 \(K=1\)와 같이 iteration 횟수를 지정하는 것이다. 그리고 maximum likelihood를 통해서 \(\hat{\beta}\)를 학습한 후, \(Q^{\beta}\)를 평가하기 위해서 policy evaluation step을 학습하게 된다. 그리고 나서 \(\pi_1\)을 학습하기 위해 아래에서 언급할 policy improvement operator를 사용한다. 핵심은 이 방식을 취하게 되면 off-policy evaluation을 완전히 피할 수 있다는 것이다.

Multi-step

Multi-step 방식을 취한다는 것은 결국 \(K \gt 1\)로 정의한다는 것이다. 이렇게 되면 evaluation operator는 \(\mathcal{D}_N\)이 behavior policy \(\beta\)에 의해서 수집되기 때문에 off-policy 환경에서 평가하게 된다. 그런데 \(K \gt 2\)인 환경에서의 evaluation 단계에서는 \(\pi_{k-1} \neq \beta\)인 policy에 대해서 평가하는 과정이 필요하게 된다. 매 iteration 마다 policy evaluation과 improvement 단계에서 수렴하도록 학습된다.

Iterative actor-critic

actor-critic 방식은 어떻게 보면 multi-step 알고리즘과 유사하지만, 매 iteration마다 수렴되도록 학습되도록 하지 않고 (애초의 목적이 다름), 훨씬 더 큰 \(K\)값을 사용한다. 여기서 iteration 마다 Q estimate를 update할 때 gradient를 취하는 과정과 policy improve하기 위한 gradient step 과정으로 구성되어 있다. 모든 evaluation과 improvement에 사용되는 operator마다 gradient를 바탕으로 계산되기 때문에 이 알고리즘은 multi-step에서 사용했던 것처럼 동일한 evaluation operator와 improvement operator를 적용할 수 있다.

Policy Evaluation operator (\(\varepsilon\))

이전에 진행된 continuous state와 action 문제와 관련된 연구처럼 이 논문에서도 단순한 fitted Q evaluation을 사용했는데, 이 방식은 DDPG처럼 target network을 사용하는 TD 형태의 학습을 통해서 최적화된다. 하지만 다른 논문처럼 double Q learning이나 Q ensemble 방식을 사용하지 않았다. one-step과 multi-step 실험시 매 iteration마다 수렴하기 위한 evaluation procedure를 학습했고, iterative 알고리즘에서는 매 iteration마다 단일 stochastic gradient step 을 밟았다. 

Policy Improvement operators (\(\mathcal{I}\))

template를 통해서 실험을 비교하기 위해서는 policy improvement operator \(\mathcal{I}\)를 지정해줘야 한다. 아래와 같은 operator를 구분지었는데, 각 operator는 behavior policy에서 hyperparameter를 통해 조절되면서 차이를 두었다.

• Behavior Cloning
  : 가장 단순한 basline은 단순히 \(\hat{\beta}\)를 새로운 정책 \(\pi\)로 반환하는 것이다. 다른 policy improvement operator는 잘 동작한다면 적어도 이 baseline보다는 잘 동작할 것이다.
• Constrained policy update
  : BCQ나 SPIBB와 같은 알고리즘은 action이 data (혹은 behavior) 의 support안에 놓이도록 policy update에 제한을 둔다. 실험을 위해서 논문에서는 Easy BCQ라고 하는 "perturbation network" 부분을 뺀 간단한 BCQ를 구현했다. 여기에서 \(\hat{\beta}\)에서 \(M\)개의 샘플을 추출한 새로운 정책 \(\hat{\pi}^M_k\)을 정의했고, \(\hat{Q}^{\beta}\)에 따라 가장 높은 값을 가지는 것을 활용했다. 공식으로 보면 다음과 같다.
  $$ \hat{\pi}^M_k (a \vert s ) = \mathbb{1}[a = \arg \max_{a_j}\{\hat{Q}^{\pi_{k-1}}(s, a_j) : a_j \sim \pi_{k-1}(\cdot \vert s), \quad 1 \ge j \ge M\}] $$
• Regularized policy updates
  : 흔히 적용할 수 있는 아이디어 중 하나는 behavior policy로 향하게끔 regularize를 취하는 것이다. 일반적인 Divergence \(D\)가 있을 때, 이에 대한 regularized objective를 최대화시키는 알고리즘을 정의할 수 있다.
  $$ \hat{\pi}^{\alpha}_k = \arg \max_{\pi} \sum_i \mathbb{E}_{a \sim \pi \vert s} [\hat{Q}^{\pi_{k-1}}(s_i, a)] - \alpha D (\hat{\beta}(\cdot \vert s_i), \pi(\cdot \vert s_i)) $$
  논문에서는 reverse KL divergence, \(KL(\pi(\cdot \vert s_i) \Vert \hat{\beta}(\cdot \vert s_i))\)를 사용했다. 참고로 reverse KL divergence를 계산하기 위해서, 우선 \(\pi(\cdot \vert s_i)\)로부터 샘플을 뽑은 후, divergence를 계산하기 위해서 해당 샘플에 대한 density estimate \(\hat{\beta}\)를 이용했다. 이렇게 하면 regularization이 \(\pi\)가 \(\beta\)의 support안에 존재하도록 할 수 있고, 이 방식은 \(\beta\)를 따라가도록 \(\pi\)에게 가산점을 주는 방식과는 다른 부분이다.
• Variants of imitation learning
  : policy improvement를 더 잘 할 수 있도록 관찰된 action에 대해서 필터링을 가하거나 가중치를 부여하는 식의 imitation learning을 수정하는 방식도 있다. 이 중 논문에서 사용한 가중치를 가하는 것은 관찰된 action에 대해서 exponentiated advantage estimate를 가중치로 사용한 것이다.
  $$ \hat{\pi}^{\tau}_k = \arg \max_\pi \sum_i exp (\tau ( \hat{Q}^{\pi_{k-1}}(s_i, a_i) - \hat{V}(s_i))) \log \pi (a_i \vert s_i ) $$

이렇게 다양한 조합을 통해서 알고리즘 template을 통한 차이 (one-step, multi-step or iterative), 그리고 improvement operator의 차이 (Easy BCQ, reverse KL regularization, or exponentially weighted imitation)를 두고 성능을 비교했다.

iterative algorithm이 잘 동작하지 않는 이유

앞에서 언급했다시피 대부분의 실험에서 one-step rl 이 multi-step이나 iterative 알고리즘보다 성능이 잘 나왔는데, 이 때의 학습 곡선을 통해서 iterative algorithm은 불안정성을 극복하기 위해서 정교한 regularization 기법의 적용이 필요하다는 것을 알았다. 그리고 이런 불안정성을 야기하는 원인으로 distribution shift와 iterative error exploitation 이 두가지를 꼽을 수 있다.

Distribution shift는 현재의 policy 를 평가하는 단계에서 고정된 dataset내에서 뽑을 샘플을 선택하게 되는데 이 크기를 줄임으로써 평가가 잘못되는 원인을 제공한다. 그래서 이전의 offline rl 논문에서도 이와 관련된 연구가 진행되고 아래에서도 언급될 예정이다. Iterative error exploitation은 dataset으로부터 추정된 Q estimate에 대해서 policy를 반복적으로 최적화하면서 내제되어 있던 error가 발산하는 현상을 말하는데, 이를 통해서 매 단계마다 overestimation 현상을 야기하는 bias를 만들게 된다. (어떻게 보면 지도학습에서 training error가 test error 보다 낮게 형성되면서 bias가 생기는 것과 유사한 것이라고 보면 좋을 듯 하다) 더불어 이렇게 dataset을 반복적으로 재사용하면서 매 step마다 training을 warm-start하기 위해 이전의 Q estimate를 사용하게 되면, 이전의 error들이 다음 단계에서 증폭되는 현상이 나타난다. 이런 현상을 특히 multi-step이나 iterative algorithm에서 잘 드러난다.

Distribution shift

off-policy evaluation을 사용하는 모든 알고리즘은 반드시 evaluation step에서 distribution shift를 겪게 된다. 이때 behavior policy와 다른 policy에 대해서 evaluation을 수행하게 되면 검증에 필요한 효과적인 sample size를 줄이게 되고, 결국 estimate에 대한 variance를 높이는 결과를 보여준다. 참고로 논문에서 언급하는 distribution shift란 behavior distribution ( dataset 상에서의 state-action 쌍에 대한 분포)와 evaluation distribution (실제 환경에서 평가할 policy \(\pi\)에 의해서 도출된 state-action 쌍의 분포)간의 차이(shift)를 나타낸다.

Iterative error exploitation

아무리 dataset내에 다양성이 풍부하고, 매 iteration마다 Q function을 처음부터 학습시켜도 off-policy evaluation을 사용하는 알고리즘은 distribution shift문제를 겪는다. 그런데 실제로 적용해보면 다른 문제도 존재한다. iterative algorithm은 Q function을 추정하면서 policy들간에 최적화하는 과정을 반복하게 되고, 이때 같은 데이터를 사용해서 Q function을 다시 추정하고, 다시 재 추정을 warm-start하기 위해서 이전 단계의 Q function을 재활용하는 과정을 반복하게 된다. 이 과정에서 iterative error exploitation이라고 하는 문제가 발생하는 요인이 각 단계별로 생긴다.

간단히 말하면 iterative error exploitation은 policy improvement 단계에서 \(\pi_i\)가 overestimate된 action을 주로 선택하면서 나타난다. 그리고 이 때의 overestimation이 policy evaluation 단계에서 Dynamic Programming에 의해 전파가 되는데, 이 것을 공식으로 표현하면 다음과 같이 설명할 수 있다. \(\pi\)에 의해서 뽑은 \(s, a\)마다 behavior policy \(\beta\)에 의해서 모인 고정된 dataset에 기반한 bellman error \(\epsilon^{\pi}_{\beta}(s, a)\)가 생기는 것이다.

$$ \hat{Q}^{\pi}(s, a) = r(s, a) \gamma \mathbb{E}_{s' \vert s, a \\ a' \sim \pi \vert s'} [ \hat{Q}^{\pi}(s', a')] + \epsilon^{\pi}_{\beta}(s, a) $$

직관적으로 생각해보면 bellman error \(\epsilon^{\pi}_{\beta}\)는 \(\beta\)에 의해서 모인 dataset에 state-action이 잘 안 모일수록 더 커질 것이다. 이 에러는 sample size가 유한하던, function approximation error에 의해서든 나타나는 error로 표현될 것이다. 문제는 이렇게 발생하는 error인 \(\epsilon^{\pi}_{\beta}\)가 다른 \(\pi\)와 상관관계가 크기 때문인데, 쉽게 설명하기 위해서 고정된 offline dataset에서 추정된 모든 policy \(\pi\)에 대해서는 \(\epsilon^{\pi}_{\beta}\)가 동일하다고 가정하고, 대신 \(\epsilon_{\beta}\)라고 표현하겠다. 그러면 이제 error는 policy와 관련성이 없어지기 때문에 error를 참의 reward를 모호하게 만드는 보조(auxiliary) reward로 바라볼 수 있게 된다. 

$$ \hat{Q}^{\pi}(s, a) = Q^{\pi}(s, a) + \tilde{Q}^{\pi}_{\beta}(s, a), \qquad \tilde{Q}^{\pi}_{\beta}(s, a) := \mathbb{E}_{\pi \vert s_0, a_0 = s, a} \Big[ \sum_{t=0}^{\infty} \gamma^t \epsilon_{\beta}(s_t, a_t) \Big] $$

사실 error가 주로 dataset에 의해서 나온다고 생각했었기 때문에 이러한 가정이 어느정도 합리적이긴 하다. 그리고 이전의 Q function이 현재의 Q function에 대해서 warm-start하기 위해서 사용되는 동안 approximation error는 자동으로 각 단계를 거치게 된다.

이제 문제에 대해서 정의해보자면, 우선 dataset이 존재한 상황에서는 \(\epsilon_{\beta}\)는 고정이 되어 있을 것이고, dataset의 support상에서 멀어질수록 더 큰 값을 가지게 될 것이다. 이 때 매 step마다 \(\pi_i\)는 \(\epsilon_{\beta}\)가 더 커지게끔 될 것이고, \(\beta\)에서 멀어질수록, \(\tilde{Q}^{\pi_i}_{\beta}\)는 \(Q^{\pi_i}\)에 상대적으로 더 큰 값으로 증가하게 된다. \(Q^{\pi_i}\)가 \(Q^{\beta}\)보다 더 학습에 좋은 정보를 제공하더라도 해당 정보는 \(\tilde{Q}^{\pi_i}_{\beta}\)에 의해서 잠식된다. 반면 \(\tilde{Q}^{\beta}_{\beta}\)는 크기가 작기 때문에, one-step 알고리즘이 해당 error에 대해서 robust하다고 볼 수 있다.

물론 환경 중에는 noisy한 state들이 많이 존재할 것이고, 이 중 일부는 overestimate될 수 있다. 매 step마다 dataset이 재사용되기 때문에 이런 overestimation은 여전히 남아서 state space상으로 전파되고 결국 iterative error exploitation을 발생시킨다. 

When are multiple steps useful?

항상 one-step algorithm이 좋은 성능을 보여주는 것은 아니다. 실제로 실험을 했을때, 임의로 수집된 dataset에서는 확실히 one-step 알고리즘이 multi-step이나 iterative 알고리즘에 비해 좋은 성능을 보여주었지만, 그렇다고 이 결과가 one-step이 multi-step에 비해서 뛰어난 성능을 보여준다는 증거가 될 수 없기 때문에 policy improvement 단계에서 multiple step을 수행할 경우 이점을 얻을 수 있는 경우에 대해서 언급하고자 한다.

이전 섹션에서 확인했다시피 multi-step이나 iterative 알고리즘은 estimation error가 전파되면서 문제가 발생한다. 이런 현상은 특히 noise가 있거다 고차원의 환경에서 문제가 도드라진다. multi-step 알고리즘은 one-step 알고리즘보다 noise를 훨씬 더 넓게 전파하기도 하지만, 한편으로는 정보를 전파하는 측면도 존재한다. 그래서 만약 noise의 크기를 줄일 수 있는 충분한 coverage가 존재한다면, 이렇게 정보가 전파되는 현상은 이점이 될 수도 있다. D4RL 상에서 진행한 실험에서는 one-step 알고리즘이 더 좋을 만큼 error가 클 때 발생할 수 있는 tradeoff에 대한 측면도 보여준다.

개인 생각

논문을 보면 강화학습 논문답지 않게 증명도 많지 않고, 대부분의 서술이 empirical evidence, 즉 실제 실험을 통해서 얻은 insight를 분석하는 과정으로 되어 있다. 그래서 논문이 어렵지 않고, offline RL에서 Off-policy Evaluation을 하면서 가지는 error에 대해 잘 도출하고, 이 중 iterative error exploitation을 줄이기 위한 간단한 baseline algorithm (논문에서는 one-step algorithm으로 지칭하고 있다)를 소개했다. 제목에 off-policy evaluation을 없다고는 했는데, 사실 딱 한번 취함으로써 error가 커지는 것을 피했다 정도로 이해했다. (참고로 나중에 소개할 논문인 IQL (Implicit Q learning)이 이 논문과 같은 시점에 발표되었는데, 잘은 기억나지 않지만 이 논문의 잘못된 점에 대해서 서술되어 있는듯 하다 뭐 아무튼...)

아래 링크는 NeurIPS 에서의 OpenReview 결과

Offline RL Without Off-Policy Evaluation