[논문 정리] VariBAD : A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning

Introduction

강화학습 (RL)은 일반적으로 보상 함수와 전이 함수가 알려지지 않은 MDP에서 기대 수익을 최대화하는 최적의 정책을 찾는 것을 목표로 한다. 이런 함수들이 주어졌다면 이론적으로 환경과 상호작용 없이 최적 정책을 계산이 가능하다.

 

그와 대조적으로 환경이 알려지지 않은 상황에서의 학습은 exploration과 exploitation의 trade-off를 요구한다. 이런 trade-off를 조절하는 것은 학습 중 기대 수익을 최대화하는 데 핵심을 담당한다.

 

Bayes-최적 정책은 이런 trade-off를 최적화하고 환경 상태뿐만 아니라 현재 MDP에 대한 에이전트 자신의 불확실성에도 조건화된 행동을 수행한다. 이런 정책은 BAMDP 프레임워크를 통해 정의할 수 있으며 여기서 에이전트는 가능한 환경들에 대한 확률적 신념 분포를 유지한다.

 

Bayes-최적 에이전트는 BAMDP의 기대 수익을 최대화하며 기대 수익을 높이는 데 도움이 되는 경우에만 불확실성을 줄이기 위한 데이터를 체계적으로 수집한다. 그 성능은 주어진 MDP에 대한 완전한 사전 지식이 있을 때 가능한 최적 정책의 성능을 상한선으로 하며 이 경우에는 탐색 없이도 최적의 행동을 취할 수 있다.

 

그러나 BAMDP에서 계획을 수행하는 것(증강된 상태에 조건화된 Bayes-최적 정책을 계산하는 것)은 대부분의 문제에 대해 계산적으로 불가능하다. 이러한 복잡성을 피하기 위해 일반적으로 사용하는 대안은 사후 샘플링이다. 여기서 에이전트는 주기적으로 사후 분포로부터 하나의 가설 MDP에서 샘플링을 진행하고 해당 MDP에 대한 최적 정책을 일정 기간 동안 따르는 방식이다. 계산 효율성은 높지만 탐색 효율성은 Bayes-최적 정책에 비해 크게 떨어질 수 있다.

그림 1. 4가지 다른 탐색 전략의 비교

 

그림 1의 예시를 고려해보자:

• 에이전트가 회색 지역에서 모르는 목표 지점으로 가야함
• 에이전트는 가능한 목표 위치들에 대해 균일한 확률을 부여
• 나머지 셀에는 확률 0을 설정하여 사후 분포를 유지
• Bayes-최적 전략은 사후 분포에서 가능한 지점들을 탐색
• 사후 분포 샘플링은 가능한 지점을 샘플링하고 거기에 이동한 뒤 목표가 없으면 분포를 업데이트하고 다시 샘플링하여 반복

이 예시는 Bayes-최적 정책이 사후 분포 샘플링보다 훨씬 더 효율적으로 탐색할 수 있음을 보인다. 핵심 도전 과제는 사후 분포 샘플링의 계산 효율성은 유지하면서 Bayes-최적에 가까운 정책을 학습하는 것이다.

 

본 논문에서는 Bayesian RL에서의 아이디어와 근사 변분 추론 그리고 메타러닝을 결합하여 주어진 MDP 분포에 따라 보지 못한 환경에서 전략적으로 탐색하고 기대 수익을 최대화할 수 있는 에이전트를 제안한다.

 

우리는 보지 못한 작업에 대해 근사적인 추론을 진행하여 메타러닝을 할 수 있는 variational Bayes-Adaptive Deep RL(variBAD)를 제안한다. MDP에 대한 분포가 주어졌을 때 우리는 단일 MDP $M$을 학습된 저차원 확률적 잠재 변수 $m$으로 표현하며 다음 두 구성요소를 함께 메타 학습한다:

• 변분 오토 인코더 : 환경과의 상호작용을 통해 수집된 경험을 바탕으로 새로운 작업에 대한 잠재 변수 $m$의 사후 분포를 추론
• 정책 : 이 MDP 임베딩의 사후 분포에 조건화되어, 작업에 대한 불확실성 하에서 탐색과 이용 사이의 trade-off를 학습

우리는 gridworld 환경과 meta-RL에서 널리 사용되는 MuJoCo 도메인에서 실험을 수행하였고, variBAD가 기존 메타러닝 방법들보다 테스트 시점에서 더 우수한 탐색 행동과 학습 중 더 높은 수익을 달성함을 확인하였다.

---

Background

우리는 Markov decision process(MDP)를 다음과 같은 튜플로 정의한다:

$$M = (\mathcal{S}, \mathcal{A}, R, T, T_0, \gamma, H)$$

• $\mathcal{S}$ : 상태의 집합
• $\mathcal{A}$ : 액션의 집합
• $R(r_{t+1}|s_t,a_t,s_{t+1})$ : 보상 함수
• $T(s_{t+1}|s_t,a_t)$ : 전이 함수
• $T_0(s_0)$ : 초기 상태 분포
• $\gamma$ : discount factor
• $H$ : 에피소드 길이

표준 RL 설정에서는 아래와 같은 정책 $\pi$를 학습한다.

$$\mathcal{J}(\pi) = \mathbb{E}_{T_0,T,\pi}\left[\Sigma^{H-1}_{t=0} \gamma^t R(r_{t+1}|s_t,a_t,s_{t+1})\right]$$

그러나 우리는 멀티태스크 메타러닝 설정을 고려할 것이다.

Training Setup

우리는 표준 메타러닝 설정을 따른다. 이 설정에서는 MDP 분포 $p(M)$가 주어지고 메타 학습동안 이 분포로부터 MDP 샘플 $M_i \sim p(M)$을 반복적으로 선택해 학습한다. 각 MDP는 아래와 같은 튜플로 정의된다.

$$M_i = (\mathcal{S}, \mathcal{A}, R_i, T_i, T_{i,0}, \gamma, H)$$

작업에 걸쳐 보상과 전이 함수는 다양하지만 일정 구조를 공유한다. 여기서 인덱스 $i$는 보지 못한 작업 표현을 표현하거나 작업 ID를 나타낸다. $p(M)$에서 MDP를 샘플링하는 것은 분포 $p(R,T)$에서 보상과 전이 함수를 샘플링하는 것으로 생각될 수 있다.

 

메타 학습 단계에서는 여러 작업들이 반복적으로 샘플링되며, 각 작업마다 간단한 학습 절차가 수행된다. 목적은 학습하는 법을 학습하는 것이다.

 

테스트 시점에는 분포에서 샘플링된 새로운 작업들에 대해 에이전트를 평가한다. 이때 평가는 에이전트가 학습을 진행하면서 얻은 평균 수익을 기준으로 이뤄진다. 이때 좋은 성능을 위해서는 두 가지가 필요하다:

• 유사한 이전 작업에서 얻은 사전 지식을 활용하는 능력
• 작업에 대한 불확실성을 추론하고 이를 바탕으로 탐색과 활용 사이의 trade-off를 조절할 수 있는 정책

이런 과제를 해결하기 위해 우리는 메타러닝과 Bayesian 강화학습의 아이디어를 결합한다.

---

 

Bayesian Reinforcement Learning

MDP가 알려지지 않았을 때 최적의 의사 결정은 탐색과 활용의 trade off를 조절해야 한다. 이론적으로는 이러한 문제를 BAMDP로 공식화하여 Bayes-최적 저책을 구하는 방식으로 해결할 수 있다.

 

Bayesian 강화학습의 관점에서는 보상 함수와 전이 함수가 사전 분포 $b_0=p(R,T)$를 따른다고 가정한다. 에이전트가 실제 보상과 전이 함수를 알 수 없기 때문에 경험 $\tau_{:t} = \{s_0,a_0,r_1,s_1,a_1,\dots,s_t\}$에 기반한 사후 분포 $b_t(R,T) = p(R,T|\tau_{:t})$를 유지한다. 이런 분포는 보통 모델 파라미터들에 대한 확률 분포 형태로 표현된다.

 

에이전트가 작업에 대한 불확실성을 의사 결정에 반영할 수 있도록 하기 위해 이 belief는 상태에 결합되어 하이퍼 상태로 확장된다.

$$s_t^+ \in \mathcal{S}^+ = \mathcal{S} \times \mathcal{B}$$

여기서 $\mathcal{B}$는 믿음 공간이다.

 

이 하이퍼 상태에 대한 전이는 다음과 같이 정의된다:

$$\begin{align}T^+(s^+_{t+1}|s_t^+, a_t, r_t) &= T^+(s_{t+1},b_{t+1}|s_t, a_t, r_t, b_t)\\ &= T^+(s_{t+1}|s_t,a_t,b_T)T^+(b_{t+1}|s_t,a_t,r_t,b_t,s_{t+1})\\ &= \mathbb{E}_{b_t}[T(s_{t+1}|s_t,a_t)]\delta(b_{t+1} = p(R,T|\tau_{:t+1}))\end{align}$$

• 새로운 환경 상태 $s_{t+1}$ : 현재의 belief $b_t$ 하에서 전이 함수에 대한 기대값
• belief $b_{t+1}$ : Bayes 규칙에 따라 결정론적으로 업데이트

보상 함수 역시 하이퍼 상태에 대해 다음과 같이 사후 분포에 따른 기대값으로 정의된다:

$$R^+(s^+_t, a_t, s^+_{t+1}) = R^+(s_t,b_t,a_t,s_{t+1},b_{t+1}) = \mathbb{E}_{b_{t+1}}[R(s_t,a_t,s_{t+1})]$$

 

이 결과로 우리는 BAMDP $M^+ = (\mathcal{S}^+, \mathcal{A}, R^+, T^+, T^+_0, \gamma, H^+)$를 얻게 된다. 이는 belief MDP의 특수한 경우로 볼 수 있다. 일반적인 belief MDP에서는 시간에 따라 변화할 수 있는 히든 상태에 대해 belief를 유지하지만 BAMDP에서는 보상 함수와 전이 함수에 대한 belief를 유지하며 이는 각 작업마다 고정된 값이다.

 

이제 에이전트의 목적은 BADMP에서의 기대 수익을 최대화하는 것이다:

$$\mathcal{J}^+(\pi) = \mathbb{E}_{b_0, T^+_0, T^+, \pi} \left[\sum^{H^+ - 1}_{t=0} \gamma^t R^+(r_{t+1}|s_t^+,a_t,s^+_{t+1})\right]$$

즉, 초기에는 환경에 대한 정보가 없는 상태에서 주어진 시간 안에 기대 수익을 최대화하는 것이다.

 

위의 목적 함수는 Bayes-최적 정책에 의해 최대화되며, 이 정책은 탐색과 활용을 자동적으로 균형 있게 조절한다. 즉, 에이전트는 오직 기대 수익을 높이는 데 도움이 될 때만 탐색을 수행한다.

 

BAMDP를 푸는 것에는 여러가지 문제점들이 있다:

• 정확한 보상 모델과 전이 모델의 정확한 파라미터를 모름
• belief 업데이트는 계산적으로 어려운 경우가 많음
• 올바른 사후 분포로부터 belief 공간을 계획하는 것은이 복잡함

이런 문제를 해결하기 위해 우리는 보상 함수와 전이 함수에 대한 추론, 그리고 belief 기반 정책 학습을 동시에 수행할 수 있는 방법을 제안한다. 이 방법은 모르는 MDP에서 추론을 수행하는 방법과 belief를 활용해 기대 수익을 최대화하는 정책을 메타러닝을 통해 end-to-end로 학습한다.

---

Bayes-Adaptive Deep RL via Meta-Learning

여기서는 variBAD를 소개하고 앞서 언급된 도전 과제들을 우리가 어떻게 해결하는지를 설명한다. 우리는 먼저 보상 함수와 전이 함수, 그리고 이에 대한 사후 확률 분포를 어떻게 표현할 수 있는지 설명한 뒤 주어진 작업에 대해 근사적인 변분 추론을 메타러닝하는 방법을 살펴보고 마지막으로 이를 모두 통합하여 전체 학습 목표를 구성한다.

 

표준적인 메타러닝 설정에서는 각 MDP마다 고유한 보상 및 전이 함수가 존재하지만 이들은 모두 분포 $p(\mathcal{M})$ 내에서 일정한 구조를 공유한다고 가정한다. 각 작업은 알려지지 않은 task description 또는 task DI $i$로 구분되지만,이런 $i$에 직접 접근할 수는 없다. 따라서 우리는 이 $i$를 직접 다루는 대신, 이를 나타내는 확률적 잠재 변수 $m_i$를 학습하여 표현한다. 즉, 주어진 MDP $\mathcal{M}_i$에 대해 다음과 같이 표현할 수 있다:

$$\begin{align}R_i(r_{t+1}|s_t,a_t,s_{t+1}) &\approx R(r_{t+1}|s_t,a_t,s_{t+1};m_i), \\ T_i(s_{t+1}|s_t,a_t) &\approx T(s_{t+1}|s_t,a_t;m_i),\end{align}$$

여기서 $R$과 $T$는 작업간의 공유가 된다. 우리는 실제 작업 표현에 접근할 수 없기 때문에 에이전트의 경험을 수집하여 $m_i$를 추론한다:

$$\tau_{:t}^{(i)} = (s_0, a_0, r_1, s_1, a_1, r_2, \dots, s_{t-1}, a_{t-1}, r_t, s_t)$$

우리의 목표는 MDP들에 대한 분포를 학습하고 환경에 대한 사후적인 정보가 주어졌을 때 최적의 행동을 선택하는 것이다. 앞서 설명한 재구성에 따르면 이제는 전이 함수와 보상 함수 전체를 직접 다루는 대신 잠재 임베딩 벡터 $m$에 대해 추론하는 것만으로도 충분하다.

Approximate Inference

실제 사후 확률을 계산하는 것은 일반적으로 불가능하다. 이는 우리가 MDP의 전체 정보에 접근할 수 없고 모든 가능한 작업들에 대해 적분을 수행하는 것은 계산적으로 매우 비효율적이기 때문이다. 따라서 우리는 두 가지를 함께 학습해야 한다.

1. 환경 모델
   • $p_\theta(\tau_{:H^+}|a_{:H^+-1})$
   • 에이전트가 취한 행동 시퀀스에 조건화된 trajectory 분포를 출력하며, 파라미터 $\theta$로 표현
2. 추론 네트워크
   • $q_\phi(m|\tau_{:t})$
   • 현재까지의 경험 $\tau_{:t}$를 입력으로 받아 잠재 임베딩 $m$의 분포를 추론, 파라미터 $\phi$로 표현

정책은 MDP의 구성요소가 아니기 때문에, 환경 모델은 행동 시퀀스가 주어졌을 때만 trajectory 분포를 생성할 수 있다. 즉, trajectory 분포는 현재 사용중인 정책 $\pi_\theta$에 조건부로 정의된다. 따라서 특정 시점 $t$에서 우리의 모델 학습 목표는 아래의 식을 최대화하는 것이다:

$$\mathbb{E}_{\rho(M,\tau_{:H^+})}[\log p_\theta(\tau_{:H^+}|a_{:H^+-1})],$$

• $\rho(M,\tau_{:H^+})$ : 정책에 의해 유도된 trajectory 분포
• 표기를 간단히 하기 위해 행동을 생략하고 오직 상태-보상 시퀀스로만 trajectory를 표현
• 이후부터는 $a_{:H^+ - 1}$에 대한 조건부 표기도 생략

위의 식을 최적화하는 것은 계산적으로 불가능하기 때문에 우리는 대신 근사 사후 분포 $q_\phi(m|\tau_{:t})$를 활용하여 추론 가능한 하한을 최적화한다. 이 lower bound는 몬테 카를로 샘플링을 통해 근사적으로 계산 가능하다:

$$\begin{align}\mathbb{E}_{\rho(M,\tau_{:H^+})}[\log p_\theta(\tau_{:H^+})] &\geq \mathbb{E}_\rho [\mathbb{E}_{q_\phi(m|\tau_{:t})}[\log p_\theta(\tau_{:H^+}|m)] - KL(q_\phi(m|\tau_{:t}) \parallel p_\theta(m))] \\ &= ELBO_t\end{align}$$

• $\mathbb{E}_q[\log p(\tau_{:H^+}|m)]$ : 재구성 손실(잠재 변수가 주어졌을 때 trajectory를 얼마나 잘 설명하는지)
• $p(\tau_{:t}|m)$ : 디코더
• $KL(q(m|\tau_{:t})\parallel p_\theta(m))$ : 사후 분포 $q_\phi$와 사전 분포 $p_\theta(m)$의 KL-divergence
• 잠재 변수 $m$에 대한 분포가 사전 지식이랑 얼마나 다른지를 측정
• 사전 분포 $p(m)$은 매 스텝 $t$에서 이전 스텝의 posterior $q(m|\tau_{:t-1})$으로 설정
• 처음에는 정규 분포로 초기화 $q(m) = \mathcal{N}(0, I)$

위의 근사 계산 식은 그림 2에서 확인 가능하다. 에이전트가 $t$에 도달했을 때 우리는 과거 trajectory $\tau_{:t}$를 인코딩하여 현재 posterior $q(m|\tau_{:t})$를 추론한다. 그런 다음 우리는 전체 trajectory $\tau_{:H^+}$- 즉, 미래를 포함한 전체 시퀀스를 디코더를 통해 복원하도록 학습한다. 즉, $\mathbb{E}_q[p(\tau_{:H^+}|m)]$을 얻는다. 과거와 미래를 같이 디코딩하도록 학습시키는 이유는 variBAD가 과거 관찰만을 보고 아직 보지 못한 미래 상태에 대해 올바른 추론을 하는 방법을 배우도록 하기 위함이다.

 

재구성 식 $\log p(\tau_{:H^+}|m)$는 아래와 같이 분해된다:

$$\begin{align}\log p(\tau_{:H^+}|m,a_{:H^+ - 1}) &= \log p((s_0,r_0,\dots,s_{t-1},r_{t-1},s_t)|m,a_{:H^+ - 1}) \\ &= \log p(s_0|m) + \sum_{i=0}^{H^+ - 1}[\log p(s_{i+1}|s_i,a_i,m) + \log p(r_{i+1}|s_i,a_i,s_{i+1},m)]\end{align}$$

• $p(s_0,m)$ : 시작 상태 분포 $T'_0$
• $p(s_{i+1}|s_i,a_i;m)$ : 전이 확률 $T'$
• $p(r_{i+1}|s_t,a_t,s_{i+1};m)$ : 보상 분포 $R'$

---

Training Objective

우리는 이제 잠재 분포로부터 task 임베딩, 정책, 보상 함수와 전이함수를 학습할 수 있다. 우리는 아래 3가지의 딥 뉴럴 네트워크를 사용했다:

1. 인코더 $q_\phi(m|\tau_{:t})$
   • 과거의 trajectory를 입력으로 받아 task 임베딩 $m$에 대한 확률 분포를 생성
   • 파라미터 : $\phi$
2. 전이/보상 모델
   
   • 전이 함수 : $T' = p_\theta^T(s_{i+1}|s_i,a_i;m)$,
   • 보상 함수 : $R' = p_\theta^R(r_{i+1}|s_t,a_t,s_{i+1};m)$
3. 정책 $\pi_\psi(a_t|s_t,q_\theta(m|\tau_{:t}))$, $\psi$로 파라미터화, $\psi$에 의존

정책은 환경 상태와 분포 $m$에 모두 영향을 받는다. 이것은 BAMDP와 유사하지만 variBAD는 전이 함수와 보상 함수를 포괄하는 잠재 임베딩 $m$ 하나로 요약하여 다룬다.

• 더 추론이 간단하고 일반화에 유리
• 모든 task의 데이터를 공유된 보상/전이 모델 학습에 활용할 수 있음
• posterior 분포는 실제로 정규분포로 가정되어 평균과 표준편차 등의 파라미터로 표현

우리의 총합 목표는 아래의 식을 최대화하는 것이다:

$$\mathcal{L}(\phi,\theta,\psi) = \mathbb{E}_{p(M)}\left[\mathcal{J}(\psi,\phi) + \lambda\sum_{t=0}^{H^+}ELBO_t(\phi,\theta)\right]$$

• 기대값은 몬테 카를로 샘플에 의해 근사화
• ELBO는 재파라미터화 트릭을 사용하여 gradient를 통한 최적화가 가능
• $t=0$에서는 $q_\phi(m) = \mathcal{N}(0,I)$를 사용
• 과거의 trajectory를 RNN 네트워크에 넣어 인코딩하지만 다른 인코딩 구조도 가능

위의 식을 보면 ELBO가 모든 가능한 콘텍스트 길이 $t$에 대해 나타나는 것을 알 수 있다. 여기서 variBAD는 어떻게 온라인에서 추론하는지를 학습할 수 있다. 그리고 더 많은 시간과 데이터가 주어질 때 불확실성을 줄이는 것을 학습한다. 실제로 우리는 고정된 숫자의 ELBO 식이 계산적으로 효율성인것을 $H^+$가 클 때 확인했다.

 

위의 식은 end-to-end로 학습하며, 이는 모델 학습(VAE의 ELBO)과 RL 정책 학습 간의 손실을 조절하는 방식이다. 이러한 구조가 필요한 이유는 인코더의 파라미터 $\phi$가 모델과 정책에서 공유되기 때문입니다. 그러나 우리는 실험을 통해 다음을 발견했다. 정책의 RL loss를 인코더까지 역전파 하는 것은 실제로 거의 필요하지 않다:

• 학습 속도 향상
• VAE와 RL 사이의 loss trade-off 조정이 필요 없어짐
• 서로 다른 성질의 gradient 간 간섭을 막을 수 있음

그래서 실험에서는

• 정책과 VAE를 완전히 분리된 optimizer와 학습률로 각각 학습
• 학습에 사용하는 데이터 버퍼도 별도 구성

으로 진행했다. 테스트 시점에서는 무작위로 샘플링한 새로운 task에 대해 인코더와 정책 네트워크를 forward pass만 수행하여 정책을 실행하고 디코더는 사용하지 않았다.

---

Related Work

Meta Reinforcement Learning

대표적인 model-free meta-RL 방법은 빠른 적응을 위한 recurrent 네트워크를 활용하는 RL2이다. 

• 매 step마다 네트워크가 이전의 action과 reward 정보를 함께 입력으로 받음
• task 내에서의 학습이 네트워크의 dynamics만으로 온라인으로 일어나도록 구성

만약 variBAD에서 디코더와 VAE 목적 식을 제외한다면 variBAD는 RL2와 거의 같은 구조로 축소된다. variBAD는 아래 두 가지의 차이점이 존재한다:

• 확률적 잠재 변수 $m$
  • 작업에 대한 불확실성을 표현할 수 있는 inductive bias
  • deterministic hidden state 대신 확률적 표현
• 디코더 + reconstruction loss
  • 과거 및 미래의 전이와 보상을 재구성하도록 학습
  • task 정보를 잠재 공간에 잘 인코딩하고 보지 못한 상태에 대한 정보도 유추하도록 도와줌

또 다른 널리 사용되는 meta-RL 접근 방식은 모델의 초기값을 학습하여 테스트 시점의 소수의 gradient step만으로 빠르게 적응하도록 만드는 방법이다. 이런 방법들은 일반적으로 초기 정책이 반드시 탐색도 잘해야 한다는 점을 고려하지 않는다.

 

이런 방식들과 모델 복잡도를 비교해보면:

• MAML, ProMP : feedforward 정책만 사용하므로 상대적으로 모델이 단순하고 가벼움
• RL2, variBAD : recurrent module을 포함하여 모델 복잡도는 높지만 대신 온라인 적응이 가능

위와 같은 방법들은 일반적으로 테스트 시점에서 탐색 단계와 활용 단계를 명확히 분리하는 구조로 설계되어 있다. 이런 설계는 구조적으로 탐색-활용을 동시에 고려하기 어렵게 만들며, 따라서 variBAD에 비해 샘플 효율성이 낮은 경향이 존재한다.

---

Skill / Task Embeddings

메타 강화학습으로 스킬 임베딩이나 작업을 학습하는 것은 다양한 방법으로 사용되었다. 

• Hausman et al. : 임베딩 공간을 근사 변분 추론으로 학습, 테스트 시점에는 정책은 고정된 채, 새로운 embedder를 학습하여 기존 스킬들 간의 보간을 수행
• Arnekvist et al. : 여러 스킬에 대해 최적 Q-함수의 확률적 임베딩을 학습하고 정책은 이 임베딩 하에 실행
• Co-Reyes et al. : 저수준 스킬의 잠재 공간을 학습하고, 이 공간은 고수준 정책이 제어, VAE를 사용하여 상태 궤적을 인코딩하고 상태/행동을 디코딩하는 방식
• Zintgraf et al. : 결정론적 태스크 임베딩을 MAML과 비슷하게 학습
• Zhang et al. : 전이와 보상 모듈을 따로 학습하여 잠재 표현을 구성, 정책은 이 표현에 조건화
• Perez et al. : 잠재 변수와 함께 동적 모델을 학습하여 이를 모델 기반 제어에 활용
• Lan et al. : MAML과 비슷한 태스크 임베딩을 학습, 테스트 시점에 인코더는 업데이트되지만 정책은 고정
• Saemundsson et al. : 모델 예측 컨트롤에 사용되는 환경 모델의 임베딩을 학습

VariBAD는 위의 방법들과 임베딩을 표현하고 사용하는 방식에서 다르다. VariBAD의 임베딩은 단순히 스킬이나 태스크 ID가 아니라 task의 불확실성까지 표현하는 확률적 분포이다. 정책은 이 분포와 직접 조건화되어 있으며, 이를 통해 task 불확실성에 대한 추론과 탐색/활용의 trade-off를 온라인 상에서 자동으로 수행할 수 있게 된다.

---

Bayesian Reinforcement Learning

RL에서의 Bayesian 접근법은 불확실성을 정량화하여 행동 선택을 지원하고 사전 지식을 알고리즘에 통합할 수 있는 방법을 제공한다. Bayes-최적 정책은 탐색과 활용의 균형을 최적으로 조절하고 학습 과정 전반에서 기대 보상을 최대화한다. 이런 정책은 이론적으로 BAMDP 프레임워크로 계산될 수 있지만, 현실적으로는 작고 단순한 작업 이외에는 계산 불가능할 정도로 비실용적이다. 

 

기존의 근사 Bayesian RL 기법들은 다음에 제한된다:

• 작고 이산적인 state/action 공간
• 작고 discrete한 task 집합

그리고 다음을 요구한다:

• 사전이나 belief 업데이트 수식의 명시적 정의
• 샘플 기반의 비싼 계획 절차

VariBAD는 메타러닝과 approximate variational inference를 결합함으로써 deep RL에서도 추적 가능한 Bayes-최적 탐색 경로를 제시한다.

 

관련 연구들에는 아래와 같이 있다:

• Wingate et al. : 정책 탐색을 Markov Chain Monte Carlo 기반 추론으로 정식화하여 추상적 지식의 task 간 전이를 시도 
• Guez et al. : Bayesian planning을 위해 MCTS를 사용한 근사 Bayes-최적 행동을 위한 tractable한 샘플 기반 방법 제안
• Osband et al. : 랜덤화된 prior 기반 탐색 방법을 통해 구조화된 탐색 제안
• Gupta et al., Rakelly et al. : 잠재 변수를 사용하여 구조화된 탐색 구조를 만들고 posterior sampling과 유사한 탐색 행동이 나타나도록 함
• Kolter & Ng : posterior를 활용한 보상 보너스 방식
• Kearns & Singh, Brafman & Tennnholtz : 불확실성을 마주했을 때의 낙관적 탐색 기반 방식

BAMDP와 관련된 개념 중 하나는 Contextual MDP이다. 이 프레임워크에서는 task를 나타내는 context가 주어지며, 이 context에 따라 환경의 전이 함수와 보상 함수가 달라진다고 가정한다. 이 방법과 variBAD의 차이를 비교하면:

• contextual MDP : Bayes-최적 행동을 직접 학습하지 않음, 새로운 환경에 오래 노출되며 학습이 진행
• variBAD : 사전 경험을 통해 빠르게 posterior를 추론하고 탐색과 활용을 동시에 처리하는 Bayes-최적 정책을 메타러닝으로 근사

---

Variational Inference and Meta-Learning

현재 존재하는 많은 Bayesian RL 방법들과 variBAD의 차이점은 추론 과정을 메타학습하는 것이다. variBAD는 강화학습 문제에서 Bayes-optimal하게 행동하기 위해 필요한 추론 과정 자체를 메타러닝으로 학습한다.

---

POMDPs

최근에는 부분 관찰 마르코프 결정 과정(POMDP) 하에서 model-free 강화학습이나 계획을 위한 모델 학습을 위한 여러 딥러닝 기반 접근 방식을이 제안되었고, 이들은 근사 변분 추론 기법을 활용한다.

 

반면, variBAD는 POMDP의 특수한 경우인 BAMDP에 초점을 맞춘다. BAMDP에서는 전이 함수와 보상 함수가 숨겨진 상태를 구성하며 에이전트는 이들에 대한 belief를 유지해야 한다. 일반적인 POMDP에서는 숨겨진 상태가 매 타임스텝마다 변할 수 있는 반면 BAMDP에서는 task 자체가 고정되어 있기 때문에 숨겨진 상태 또한 task마다 고정되어 있다.

 

variBAD는 이 속성을 활용하여 시간에 따라 변하지 않는 고정된 임베딩을 학습한다. 이는 매 timestep마다 변화하는 숨겨진 상태를 추적하는 필터링 기반 접근법과 대조적이다.

---

Experiments

우리는 먼저 didactic 그리드월드 환경에서 variBAD의 특성을 실험적으로 분석한다. 실험을 통해 variBAD가 현재 task를 추론하면서 구조적이고 온라인적인 탐색을 수행함을 확인한다. 그 후, 더 복잡한 메타러닝 세팅인 MuJoCo에서 연속적인 컨트롤 작업을 진행했다. 우리는 variBAD가 첫 rollout에서 task에 적응하는 모습을 보였다.

Gridworld

이 환경에서의 과제는 5x5 격자 안에서 무작위로 선택된 목표 지점에 도달하는 것이다. 에이전트는 목표의 위치를 관찰할 수 없기 때문에, 작업에 대한 불확실성이 존재하며 탐색이 필수적이다. 목표는 시작 지점을 제외한 어느 위치든 가능하다.

 

에이전트에 대한 정보는 아래와 같다:

• 가능한 행동 : 위, 오른쪽, 아래, 왼쪽, 정지 -> 결정론적으로 실행
• 에피소드는 15스텝 후 리셋
• BAMDP 상에서는 4개의 MDP 에피소드를 포함한 60으로 horizon을 설정
• 보상은 희소하게 주어짐(목표 위치에서만 +1)
• 잠재 변수는 5차원으로 설정

그림 3은 그리드 환경에서 variBAD의 행동을 나타낸 것이다.

그림 3. 그리드월드 환경에서의 variBAD의 행동

1. 그림 3a. 테스트 시점의 행동 분석
   • 탐색은 전적으로 정책 자체의 구조화된 행동에 의해 이뤄짐
   • 파란 색상은 보상 함수로부터 유도된 사후 분포를 시각화
   • variBAD는 올바른 사전 분포를 학습하고 점진적으로 사후 분포를 업데이트
   • 이미 방문한 셀들에는 보상이 없다고 판단하며 남은 셀들에 대해 효율적으로 탐색 진행
2. 그림 3b. 보상 예측 분석
   • variBAD는 디코더가 예측하는 보상을 통해 정책이 어떤 환경을 믿는지(belief)에 대한 통찰을 얻을 수 있음
   • 3b에서는 각 셀이 보상을 받을 확률을 선으로 시각화
   • 데이터를 수집할수록 보상이 없다고 판단되는 셀이 점점 늘어남
3. 그림 3c. 잠재 공간 분석
   • 5차원 잠재 공간의 변화를 시각화
   • 목표를 찾기 전에는 잠재 변수 m의 분포가 퍼져 있음
   • 목표를 찾은 후에는 posterior가 집중되며 분산이 거의 0에 가까워짐
   • 에이전트가 이 task가 어떤 것인지 확신을 가지게 되었음을 의미

variBAD의 행동은 그림1에 나오는 Bayes-최적 정책과 매우 유사하다. 우리의 실험은 variBAD가 Bayes-optimal한 탐색/활용을 근사화하는 데 매우 효과적이며 정책이 실제로 무엇을 믿고 행동하는지를 들여다볼 수 있는 분석적 장점을 가진다는 것을 보여준다.

---

MuJoCo Continuous Control Meta-Learning Tasks

우리는 variBAD가 복잡한 메타러닝 환경에서도 잘 확장됨을 보이기 위해, meta-RL 분야에서 널리 사용되는 MuJoCO 기반 로코모션 과제에 적용해 실험을 수행했다. 구체적으로 다음 네 가지 환경을 실험에 사용했다:

1. AntDir 및 halfCheetahDir: 에이전트가 앞 또는 뒤로 달려야 하는 두 가지 task 중 하나를 수행해야 함
2. HalfCheetahVel: 에이전트가 다양한 속도로 달리는 task들을 수행해야 함
3. Walker: 에이전트의 시스템 파라미터가 무작위로 샘플링되는 환경

그림 4는 테스트 시점에서의 성능을 기존 방법들과 비교한 것이다. 완전한 비교를 위해 여러 rollout 결과를 모두 제시했지만 우리의 목적은 단 한 번의 episode 내에서 새로운 task에 적응하면서 최대 성능을 내는 것이다. 따라서 첫 번째 rollout이 가장 중요하다.

그림 4. MuJuCo에서의 5 rollout 평균 성능 결과

결과를 보면:

• variBAD와 RL2만이 단일 episode 내에서 즉각적으로 task에 적응하는 능력을 보임
• 그러나 RL2는 HalfCheetahDir 환경에서 variBAD보다 성능이 낮고, 학습 속도도 느리고 안정성이 떨어짐
• 첫 rollout에 탐색이 포함되어 있음에도 불구하고, variBAD의 성능은 oracle policy와 거의 유사

다른 메타러닝 방법들 PEARL, E-MAML, ProMP는 단일 rollout에서 보상을 최대화하도록 설계되지 않았기 때문에, variBAD와 같은 조건에서는 성능이 크게 떨어졌다. 이들 방법은 모두 새로운 task마다 훨씬 더 많은 환경 상호작용이 필요하며, 좋은 성능을 내려면 여러 episode에 걸쳐 학습해야 한다.

---

Conclusion

우리는 Bayes-최적 행동을 근사하기 위한 새로운 deep 강화학습 방법 variBAD를 제안했다. 이 방법은 메타러닝을 통해 유사한 작업들로부터 얻은 지식을 활용하고 알 수 없는 환경에서의 근사 추론을 수행할 수 있도록 설계되었다. 직관적인 그리드월드 환경에서 variBAD는 Bayes-최적 행동과 매우 유사한 방식으로 작동했고 복잡한 MuJuCo 실험들에서도 단일 에피소드 내에서의 보상 측면에서 기존 방법들을 능가하는 성능을 보였다.

 

variBAD를 기반으로 다양한 확장 연구가 가능하다:

• 디코더를 모델 예측 계획에 활용
• 디코더의 예측 오류를 분석하여 환경이 학습 분포에서 벗어났는지 활용
• 훈련 환경과 테스트 환경의 분포가 다른 경우로 확장

특히 variBAD에는 두 가지 문제가 생길 수 있다:

• 추론 과정이 잘못될 수 있음
• 정책이 변화된 posterior를 올바르게 해석하지 못할 수 있음

이런 경우에는 인코더/디코더의 추가적인 학습, 정책의 재학습 또는 명시적 planning의 도입이 필요할 수 있다.

 

https://arxiv.org/abs/1910.08348

VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning