Dream to Control: Learning Behaviors by Latent Imagination
저자: Danijar Hafner et al.
학회: ICLR
년도: 2020
Keyword
Model-based reinforcement learning, latent imagination
Main contribution
1. Learning long-horizon behaviors by latent imagination
주어진 state를 latent로 mapping하여 model-based learning을 수행함. 학습된 imagination 환경 모델을 활용하여 policy를 효율적으로 학습할 수 있음 (action model에 대한 analytic gradient)
2. Empirical performance for visual control
Vison-based 제어 환경에서 실험을 수행하였으며, 다른 model-based, model-free 방법론에 비하여 적은 데이터를 사용하여 높은 성능을 보임. 또한, computation time도 개선
Dreamer의 전체 구조
본 논문에서 제안하는 Dreamer는 (a) dynamics learning, (b) behavior learning, and (c) environment interaction으로 구성.
- Dynamics learning: 주어진 데이터를 활용하여 환경의 특성을 학습하여 World model을 생성.
- Behavior learning: 학습된 World model을 사용하여 value를 평가하고 action을 결정.
- Environment interaction: 학습된 actor로 환경에서 데이터를 sampling.
Dreamer의 학습 과정
Learning action and value model
Dreamer 구조는 actor-critic 학습 구조를 사용함. Action model (actor) 는 imagination 환경을 풀어내는 action을 추론하는 것을 목적으로 하고, value model은 imagination 환경에서 발생하는 state $s_\tau$에 대한 value를 추론하는 것을 목적으로 함.
- Action model : $a_\tau\ \sim \ q_\phi(a_\tau\ |\ s_\tau)$
- Value model : $v_\psi(s_\tau) \approx \mathbb{E}_{q(\cdot | s_\tau)}(\Sigma_{\tau=t}^{t+H} \gamma^{\tau-t}r_\tau) $
학습과정은 기본적인 actor-critic 학습 과정을 따름. Value model은 return, advantage 등 다양한 objective를 추론할 수 있지만, Dreamer는 다음과 같은 value model을 사용.
$$
V_N^k(s_\tau) = \mathbb{E}_{q_\theta, q_\phi} (\sum_{n=\tau}^{h-1}\gamma^{n-\tau}r_n + \gamma^{h-\tau}v_\phi(s_h))\ with\ h=min(\tau+k, t+H)
$$
$$
V_\lambda(s_\tau) = (1-\lambda)\sum_{n=1}^{H-1}\lambda^{n-1}V_N^n(s_\tau) + \lambda^{H-1} V_N^H(s_\tau)
$$
$V_N^k(s_\tau)$는 $s_\tau$ k-step cumulative reward에 k-step 이후 state의 value를 더한 값을 뜻함. Dreamer는 $V_N^k$를 k에 대한 exponentially weighted average를 하여 사용.
따라서, 다음과 같은 objective로 action model과 critic model을 학습함.
$$
\max_\phi \mathbb{E}_{q_\theta, q_\phi}(\sum_{tau=t}^{t+H}V_\lambda(s_\tau)), \ \ \ \min_\psi \mathbb{E}_{q_\theta, q_phi}(\sum_{\tau=t}^{t+H}\frac{1}{2}\lVert v_\phi(s_\tau)-V_\lambda(s_\tau)\rVert)
$$
기존 actor-critic 모델과는 다른 점은, actor에 대한 직접적인 gradient (본 논문에선 analytic gradient라고 설명) 를 줄 수 있다는 점인데, 이는 reward, transtion model $q_\theta$가 모두 미분 가능한 신경망으로 이루어져 있기 때문임.
위의 $V_\lambda(s_\tau)$를 구하는 식을 보면, 결국 value는 reward model와 다음 state를 추론하는 transition model에 의해 결정됨. Reward는 state에 의해 결정되고, state를 생성하는 transition model은 action, 즉 action model의 output을 input으로 받으므로, action model 까지 gradient path가 연결됨. 따라서, action model에 대해 직접적으로 value를 maximize 하도록 gradient를 계산하여 학습할 수 있음.
Learning latent dynamics
Reward prediction
Model-based에서 제안할 수 있는 가장 기본적인 모델. Transition probability와 reward function을 학습하여 이를 actor-critic 모델에 활용하는 것을 목적으로 하고 있음. 이 모델은 아래와 같이 구성됨.
- Representation model : $p_\theta(s_t\ |\ s_{t-1},\ a_{t-1},\ o_t)$
- Transition model : $q_\theta(s_t\ | s_{t-1},\ a_{t-1})$
- Reward model : $q_\theta(r_t\ |s_t)$
이런 방식은 데이터가 충분할 시에는 좋은 성능을 보이나, 데이터가 불충분하거나 reward가 sparse할때는 성능이 저하됨.
Reconstruction
본 논문에서 제안하는 Imagination 환경은 다음과 같은 네가지 모델로 구성됨.
- Representation model : $p_\theta(s_t\ |\ s_{t-1},\ a_{t-1},\ o_t)$
- observation model : $q_\theta(o_t\ |\ s_t)$
- Transition model : $q_\theta(s_t\ | s_{t-1},\ a_{t-1})$
- Reward model : $q_\theta(r_t\ |s_t)$
Representation model은 observation과 action을 continous vector인 state $s_t$ 로 encode 하고, transition model은 현재 state와 action을 사용하여 다음 state를 추론하며, reward model은 현재 state를 기반으로 reward를 추론함.
목적 함수는 varitional lower bound (ELBO) 를 사용함. 이 loss는 VAE 같은 모델을 학습시킬 때 기본적으로 사용하는 목적 함수로, encoder가 observation을 latent로 변환할 때 특정 distribution (prior)를 유지하도록 함.
$$
\mathcal{J}_{REC} = \mathbb{E}_p(\sum_t(\mathcal{J}_O^t + \mathcal{J}_R^t + \mathcal{J}_D^t) + const \ \ \ \ \mathcal{J}_O^t = ln\ q(o_t|s_t)
$$
$$
\mathcal{J}_R^t = ln\ q(r_t|s_t) \ \ \ \ \mathcal{J}_D^t = -\beta KL(p(s_t|s_{t-1}, a_{t-1}, o_t) || q(s_t|s_{t-1},a_{t-1}))
$$
$\mathcal{J}_O^t$ 는 주어진 state $s_t$에 대해서 observation $o_t$, $\mathcal{J}_R^t$ 는 주어진 state $s_t$에 대하여 reward $r_t$ reconstruction을 위한 목적 함수이고, $\mathcal{J}_D^t$는 observation이 없을 상태의 distribution (prior)에 observation을 본 후의 distribution (posterior)를 맞추는 목적 함수.
Transition model은 recurrent state space model (RSSM)을 사용하였고, image에 대한 representation learning을 위하여 convolution network를 결합함.
Constrastive estimation
하지만, observation이 이미지인 경우에는 high-dimension인 이미지를 reconstruction 해야 한다는 단점이 있음. pixel-wise한 이미지 생성을 해야하는, state를 사용하여 observation을 추론하는 observation model 대신, observation을 사용하여 state를 추론하는 state model을 사용함.
- State model : $q_\theta(s_t\ |\ o_t)$
이 모델을 학습시키기 위하여, noise contrastive estimation 방식을 사용하는데, state와 observation 간의 mutual information을 최대화 시키기 위하여, Contrastive estimation을 수행함.
$$
\mathcal{J}_{NCE} = \mathbb{E}(\sum_t(\mathcal{J}_S^t + \mathcal{J}_R^t + \mathcal{J}_D^t)) \ \ \ \ \mathcal{J}_S^t = ln\ q(s_t|o_t) - \ln(\sum_{o'}q(s_t|o'))
$$
이렇게 observation model 대신 state model을 사용하고, 목적함수를 그에 따라 변경할 수 있는 이유는, 앞서 설명했던 목적 함수의 $ \mathcal{J}_O^t + \mathcal{J}_R^t $ 부분에서 state의 constant marginal probability를 빼고 Bayes rule을 사용하면 $\mathcal{J}_{NCE}$ 가 $\mathcal{J}_{REC}$의 lower bound임이 증명되기 때문임.
실험
실험 환경
DeepMind Control Suite에 있는 20개의 visual task
실험에 사용한 비교군
- A3C (+PlaNet): PlaNet은 2018년에 나온 동저자의 Model-based 방법론으로, 대부분의 방법론이 유사하지만, action model을 사용할 때 analytic gradient를 사용하여 학습한 것이 제일 큰 차이점임.
- D4PG (+PlaNet)
- Dreamer + Reconstruction: imagination 환경의 목적 함수를 $J_{REC}$를 사용한 모델.
- Dreamer + Constrative (Constrastive estimation): imagination 환경의 목적 함수를 $J_{NCE}$를 사용한 모델.
- Dreamer + Reward only (Reward prediction): transition model, value estimation 없이 주어진 state에 따라서 reward만 추론하는 방식.
실험 성능
- 대부분의 환경에서 Dreamer를 사용한 모델이 빠른 학습 속도를 보임.
- Reconstruction을 사용하는 것이 construction을 사용하는 것 보다 높은 성능을 보임.
결론
- Latent 기반의 model-based 방법론을 활용하여 복잡한 환경인 visual control task를 데이터 효율적이고 적은 연산량으로 학습할 수 있음.
- 환경을 미분 가능한 학습 모델에 근사함으로써, 효율적인 강화학습을 수행할 수 있게 됨.
Reference
Hafner, D., Lillicrap, T., Ba, J., & Norouzi, M. (2019). Dream to Control: Learning Behaviors by Latent Imagination. http://arxiv.org/abs/1912.01603
