[CS285] 11. Model-Based Reinforcement Learning

💡 이 글은 『2024 PseudoLab 전반기 강화학습팀』으로 진행되었으며, CS285 Fall 2023를 따라 정리했습니다.

1. Basics of Model-based RL

저번 Lecture에서는 modeling이 모두 이루어진 경우, 즉 $s_ {t+1} = f(s_t, a_t)$로 deterministic하게 표현되거나, $p(s_ {t+1} \vert s_t, a_t)$로 stochastic하게 표현되는 것을 알 때 어떻게 모델을 활용하는지에 대해 알 수 있었다. 이번 Lecture부터는 본격적으로 그 관계, modeling을 어떻게 학습할지에 대해 다룬다.

가장 쉽게 생각해볼 수 있는 방법은 다음과 같다. 아직 불완전하기에 model-based RL version 0.5라고 부르겠다.

이 방법은 아주 직관적이다. Transition sample $(s, a, s^ \prime)_ i$를 얻어서 이를 supervised manner로 학습시킨다. 이렇게 얻은 dymanics model을 downstream task에 활용한다. 이 방법은 간단하고, 대부분의 부분을 모델링할 수 있는 비교적 쉬운 시스템에서는 잘 작동하지만, 그 외에는 잘 작동하지 않는다.

그 이유는 distributional shift 때문이다. Random policy로 얻은 분포 $p_ {\pi_ 0} (s_ t)$와 실제로 planning algorithm을 통해 얻어지는 분포 $p_ {\pi_ f} (s_ t)$는 다르다. 이는 imitation learning 때 생겼던 문제와 동일한 문제이다. 특히 모델이 expressive할수록(high capacity일수록) distribution에 더 잘 fit하기 때문에 문제가 더 커진다.

따라서 이를 해결하기 위해 실제로 작동하는 model-based RL version 1.0을 생각해보자.

이처럼 실제 policy $\pi_ f$를 통해 얻은 sample을 추가로 이용해 model을 학습한다. 이렇게 하면 distribution을 더 유사하게 만들 수 있기에 distributional shift 문제를 어느 정도 해결할 수 있다.

더 나아가 model-based RL version 1.5를 생각해보자. 현재 방법은 action sequence를 한 번에 정한 뒤 이를 실행하는 방법이다. 이는 mistake가 커지기 전에 교정할 수 없다는 단점이 있다. 따라서 이를 해결하기 위해 MPC(model predictive control)을 사용한다. 이는 plan sequence를 계속 고치는 방법으로, mistake가 커지기 전에 교정할 수 있기 때문에 더 robust하다. 이를 replanning이라고 한다.

대신 이러한 replanning을 너무 자주 하게 되면 전체 action sequence의 완성도가 떨어지게 되므로, 이를 trade-off해야 한다.

2. Uncertainty in Model-based RL

2.1. Uncertainty Estimation

문제는 model-based RL의 성능이 model-free RL보다 훨씬 나쁘다는 점이다. 아래 그래프를 보면, model-based RL으로 학습시킬 때보다 model-free RL로 학습시킬 때 성능이 비약적으로 상승하는 것을 알 수 있다.

왜일까? 이 문제는 overfitting과 관련이 있다. 초반 학습에서는 distribution shift를 피하기 위해 overfitting을 방지해야 하지만, 결국 학습을 진행시키면서 high capacity를 가진 모델이 필요하다. 그러나 이 trade-off를 잘 해결하지 못하기에 성능이 나쁘다. Model-free RL은 이러한 dynamics distribution을 고려할 필요가 없다.

따라서 이를 해결하기 위해 uncertainty estimation을 사용한다. 이는 model이 얼마나 확신을 가지고 있는지를 측정하는 방법이다. 이를 통해 model이 잘못 예측할 수 있는 곳으로 가지 않도록 하여 더 안전한 action을 취할 수 있다. 따라서 모델이 오로지 exploit만 하는 것을 방지하고, 충분한 explore을 통해 distribution을 신뢰 가능한 수준으로 만들 수 있다. 참고로 여기서 expected value는 pessimistic이나 optimistic value와 다르다. 따라서 정확한 측정값이 아니나 uncertainty를 측정하는 데에는 충분하다.

2.2. Uncertainty-Aware Models

쉽게 생각해볼 수 있는 방법은 model의 output으로 나온 probability distribution을 활용하는 것이다. 즉 model의 output entropy를 활용하여 얼마나 model이 확신을 가지고 있는지를 측정한다.

그러나 이는 적절하지 않다. Uncertainty에는 두 가지 타입이 있는데, statistical uncertainty와 model uncertainty이다. Statistical uncertainty는 nature uncertainty로, 주사위를 던졌을 때 나오는 숫자의 uncertainty와 같다. 반면 model uncertainty는 data가 부족해서 생기는 uncertainty로, 모델에 우리가 아직 확신을 가지지 못하는 부분이다. 따라서 model uncertainty를 측정하기 위해서는 output uncertainty만으로는 충분하지 않다. 모델은 확신을 가지고 예측하더라도 우리가 모델에 확신을 가지지 못하기 때문이다.

따라서 model uncertainty, 즉 $p(\theta \vert \mathcal{D})$를 측정하고자 한다. 그리고 실제로 $p(s_ {t+1} \vert s_t, a_t)$를 다음과 같이 예측하는 것이다.

\[\int p(s_ {t+1} \vert s_t, a_t, \theta) p(\theta \vert \mathcal{D}) d\theta\]

간단히 생각해볼 수 있는 방법은 다음과 같은 BNN(Bayesian Neural Network)를 사용하여 uncertainty를 모델링하는 것이다. BNN은 단순한 숫자 대신 distribution으로 구성되어 있으며, 이를 통해 uncertainty를 측정할 수 있다. 물론 이는 computationally expensive하다.

좀 더 간단한 방법은 bootstrap ensemble을 사용하는 것이다. 이는 bootstrap sampling을 통해 ensemble을 구성하고, 이를 통해 distribution을 근사하는 방법이다. 이는 BNN보다 computationally efficient하다.

가장 이상적인 bootstrap 방법은 independent dataset으로 각각 independent model을 훈련하는 것이다. 그러나 이는 쉽지 않다. 따라서 단순히 여러 번 훈련하는 방법을 사용한다. 그래도 SGD와 random initialization 때문에 충분히 independent하기 때문이다.

2.3. Planning with Uncertainty

이제 uncertainty를 고려해서 모델 $i$에 대한 output을 종합적으로 고려하여 다음과 같은 과정을 거친다. 이후 planning 방법은 Lecture 10에서 다룬 것과 같으므로 생략한다.

이러한 uncertainty estimation을 통해 model-based RL의 성능을 향상시킬 수 있으며, 실제 실험 결과 model-free RL을 넘는 성능을 보이기도 한다.

3. Model-based RL w/ Complex Observations

지금까지는 state space가 low-dimensional인 경우를 가정했다. 그러나 high-dimensional observation을 다루는 경우에는 어떻게 해야 할까? 가장 쉽게 접할 수 있는 예시는 image-based observation이다. 이 경우 image를 그대로 사용할 수도 있지만, 이는 computationally expensive하다. 따라서 image를 feature로 변환하여 사용하는 것이 일반적이다. 따라서 다음과 같은 latent space model을 고려하자.

여기서 observation $o_t$는 실제 image, state $s_t$는 latent space라고 생각하면 된다. 따라서 dynamics model은 latent space에서 학습되고, 추가로 observation model $p(o_t \vert s_t)$을 학습한다. 이를 통해 observation을 latent space로 변환할 수 있다. 따라서 training objective를 다음과 같이 수정해야 한다.

이때 posterior $p(s_ t, s_ {t+1} \vert o_ {1:T}, a_ {1:T})$로부터 $(s_ t, s_ {t+1})$를 샘플링하여야 하고, 따라서 posterior를 알아야 하지만 이는 intractable한다. 따라서 우리가 알 수 있는 분포 $q_ \psi (s_t, s_ {t+1} \vert o_ {1:T}, a_ {1:T})$로 근사한다. 이를 variational inference라고 한다. 두 가지 방법이 있는데, full-smoothing posterior는 정확하지만 복잡하고, single-step encoder는 간단하지만 덜 정확하다. 단순화를 위해 single-step encoder를 생각해보자.

그 중에서도 deterministic encoder를 사용하는 아주 간단한 방법만 생각해보자. 더 복잡한 내용은 Lecture 12의 variational inference에서 다룬다.

이때 $s_ t$에 $g_ \psi (o_ t)$를 대입하여 다음과 같이 objective를 수정할 수 있다. 이후 backpropagation을 통해 학습할 수 있다.

Reward model $p(r_ t \vert s_ t, a_ t)$를 추가하면 최종적인 objective는 다음과 같고, 결과적으로 알고리즘은 아래와 같다.

CS285 카테고리 내 다른 글 보러가기