Muesli LunarLander-v2 implementation - 상세구현

개요

해당 글에서는 구현된 사항들 몇몇에 대해 추가로 설명합니다.

Muesli (LunarLander-v2) 코드 : 

네트워크

네트워크는 간단한 mlp들로 구성이 되어 있고(Linear layer + ReLU activation, width 128), prediction network 및 dynamics network는 앞쪽에서 parameter sharing하고 뒤쪽에서 각 출력에 해당하는 head로 나눠지는 구조입니다. Representation network는 표현 학습이 중요한 부분이라 강화시킬 부분이 있을까 해서 skip connection형태의 연결을 추가했는데 이게 생각보다 도움이 많이 됩니다. 

출력되는 hidden state들은 매번 [-1,1] 범위 내로 min-max normalization 사용해서 스케일 조정이 되는데 이 부분이 없으면 hidden state 표현 학습되는 범위가 그때마다 천차만별이라 학습이 불안정하게 되고 NaN오류 뜨면서 학습이 중지되기도 하기 때문에 꼭 필요한 부분입니다. 

Action은 LunarLander의 경우 0,1,2,3 4개가 있는데, dynamics network에 들어갈 때 one-hot표현으로 만들어서 state 옆에 붙이는 형태로 구현을 하였습니다.

(Atari같은 경우는 CNN을 쓰니까 action을 plane으로 만들어서 stack을 하고, hidden state도 input RGB범위인 [0,1]로 스케일 조정을 하는 차이가 있음)

각 value model, reward model의 경우 (supprot size*2+1 크기의) categorical 분포를 출력합니다.



Self-play시에는 아래와 같이 동작합니다.

Learning rate scheduler (Linearly decay)

뮤즐리는 시작시부터 learning rate를 linear하게 decay시켜 학습이 끝날때 0이 되게끔 하는 스케쥴러를 사용합니다. Pytorch 공식 구현에는 linear하게 내려가는 스케쥴러가 없어서 https://github.com/cmpark0126/pytorch-polynomial-lr-decay 에서 제공하는 스케줄러를 사용했습니다. 

self.scheduler = PolynomialLRDecay(self.optimizer, max_decay_steps=4000, end_learning_rate=0.0000

이런 스케쥴러를 쓰는 이유는 뮤즐리가 advantage normalization을 사용하기 때문입니다. 학습이 진행되면서 policy가 optimal policy에 가까워져서 advantage가 작아져도 normalize된 advantage는 작지 않기 때문에 학습이 수렴하기 어려운데 이를 learning rate를 줄여가는 것으로 해결을 하는 것입니다.


Advantage normalization

Advantage를 normalization없이 사용하면 advantage가 reward scale에 따라서 같이 크기가 변동하게 되고 알고리즘 내에서 사용하는 clip에 항상 걸리거나 loss가 너무 커져 학습이 불안정하게 됩니다. 

때문에 뮤즐리 논문에서 advantage를 normalize하는 방법을 제시하고 있고 이에 따라 advantage를 나눠주면 너무 작은 경우 크게 되고, 너무 큰 경우는 줄어들기 때문에 advantage값이 크게 튀지가 않아서 안정적인 학습이 가능해집니다. 


        self.var = 0
        self.beta_product = 1.0

        beta_var = 0.99
        self.var = beta_var*self.var + (1-beta_var)*(torch.sum((G_arr_mb[:,0] - to_scalar(t_first_v_logits))**2)/16)
        self.beta_product *= beta_var
        var_hat = self.var/(1-self.beta_product)
        under = torch.sqrt(var_hat + 1e-12)

        모든 adv에 대해 adv <- adv/under





Mini-batch 구성 및 replay

Replay buffer로부터 16개의 sequence들(4개 replay * 4개 sequence)을 랜덤으로 선택해서 배치로 엮어 한번의 iteration에 사용합니다. 논문에서는 5% 정도의 online 데이터 사용을 하는데 저의 경우는 배치 사이즈가 작아서 100% 랜덤으로 사용했습니다.

배치 만들때 trajectory길이를 넘어가는 sequence가 골라지는 경우가 있으니 넘어가는 부분에는 zeros 텐서가 쓰이게 했습니다.
 

Dynamics network gradient scale

Unroll할때 나온 hidden state들에 .register_hook(lambda grad: grad * 0.5) 을 걸어줘서 gradient계산될 때 앞쪽으로 갈수록 gradient가 절반씩 줄어들게끔 해주어 앞쪽까지 과도한 gradient가 중첩되는 것을 막아줍니다.


Gradient clipping by value

이는 논문에서 제시된 내용으로 gradient를 [-1,1] 범위로 clip 해줍니다. Optimize하기 전에
nn.utils.clip_grad_value_(self.parameters(), clip_value=1.0) 를 사용하였습니다.


Target network(prior parameters) moving average update

뮤즐리에서는 업데이트되었던 네트워크의 파라미터들을 exponential moving average로 기억하고 있는 target network가 쓰이는데 이를 사용하기 위해 업데이트 된 파라미터를 target network와 적정 비율로 계속 섞어주는 동작을 합니다. 


            ## target network(prior parameters) moving average update
            alpha_target = 0.01 
            params1 = self.named_parameters()
            params2 = target.named_parameters()
            dict_params2 = dict(params2)
            for name1, param1 in params1:
                if name1 in dict_params2:
                    dict_params2[name1].data.copy_(alpha_target*param1.data + (1-alpha_target)*dict_params2[name1].data)
            target.load_state_dict(dict_params2)    




Stacking frames

Self-play 및 학습에서 observation을 한 프레임만 사용해서 계산을 하는 것보다 이전 프레임들을 함께 사용하면 성능에 유리한 부분이 있습니다. 저는 8프레임을 사용했습니다. (뮤즐리 논문에서는 16프레임, 뮤제로는 32프레임 사용)

Hyperparameters

learning rate : 0.00032 -> 0
Discount : 0.997
Unroll : 5 step

regularizer_multiplier : 5
Value loss weight : 0.25
Reward loss weight : 1.0

Variance moving average decay beta_var : 0.99
Variance offset epsilon : 1e-12
Target network update rate alpha target : 0.01



댓글

가장 많이 본 글

구글 람다(LaMDA)란? - 구글의 언어 모델

알파고 강화학습 원리

버텍스 AI란? - 구글 인공지능 플랫폼

카타고와 바둑 두어보기

뉴럴 네트워크란?

블로그 글 목록

뉴럴 네트워크를 학습시키는 방법