28.Policy Gradient 실습 1

( 참고 : Fastcampus 강의 )

[ 28. Policy Gradient 실습 1 ]

1. REINFORCE 복습

2. Import Packages

import sys; sys.path.append('..') 

import gym
import torch
import matplotlib.pyplot as plt

from src.part3.MLP import MultiLayerPerceptron as MLP
from src.part4.PolicyGradient import REINFORCE
from src.common.train_utils import EMAMeter, to_tensor

import torch.nn as nn
from torch.distributions.categorical import Categorical

3. REINFORCE 알고리즘

(1) __init__

• policy : 정책 ( 여기서 policy는 NN 모델로써, 계속 update할 것임 )
• gamma : discount rate
• opt : Adam optimizer
• _eps : to prevent numerical problems of logarithms

(2) get_action

• policy NN의 output은 “logit 형태”이다.
• 여기서 나온 logit 값을 기반으로 action을 샘플한다.

class REINFORCE(nn.Module):
    def __init__(self,policy: nn.Module,gamma: float = 1.0,lr: float = 0.0002):
        super(REINFORCE, self).__init__()
        self.policy = policy  
        self.gamma = gamma
        self.opt = torch.optim.Adam(params=self.policy.parameters(),lr=lr)
        self._eps = 1e-25

    def get_action(self, state):
        with torch.no_grad():
            logits = self.policy(state) # output : logit 형태
            action_dist = Categorical(logits=logits)
            action = action_dist.sample()  
        return action

4. Functions ( preprocess, update )

(1) _pre_process_inputs

연산의 효율성을 위해서, 들어온 episode ( S,A,R )를 reverse flip해주는 함수이다.

@staticmethod

• class에서 method를 바로 호출할 수 있음

• self를 인자로 받지 않는다
• instance 속성에 접근할 수 없음 ( 필요 없을 때 주로 사용 )
• 순수 함수(pure function)를 만들 때 사용

@staticmethod
def _pre_process_inputs(episode):
    s, a, r = episode
    s = s.flip(dims=[0]) # [num.steps x state_dim]
    a = a.flip(dims=[0]) # [num.steps]
    r = r.flip(dims=[0]) # [num.steps]
    return s, a, r

(2) update

• 매 time-step마다 update가 이루어짐 ( 실제로는 비효율적이지만, 실습을 위해서 구현 )

  ( 아래 코드를 보면, episode가 input으로 들어오고, update를 하기 위해 epsiode내의 각 time step을 (zip을 통해) 계속 뽑아냄을 알 수 있다 )

• (식 A) \(G_{t} \leftarrow \sum_{k=t+1}^{T} \gamma^{k-t-1} R_{k}\)

• (식 B) \(\theta \leftarrow \theta+\alpha \gamma^{t} G_{t} \nabla_{\theta} \ln \pi_{\theta}\left(A_{t} \mid S_{t}\right)\)

def update(self, episode):
    S, A, R = self._pre_process_inputs(episode)
    g = 0
    for s, a, r in zip(S, A, R):
        ### (식 A)
        g = r + self.gamma * g
		
        ### (식 B)
        # minimize해야하는 LOSS이므로, '-'붙이기
        action_dist = Categorical(logits=self.policy(s))
        action_prob = action_dist.probs[a]
        policy_grad_loss = - torch.log(action_prob + self._eps) * g
        
        self.opt.zero_grad()
        policy_grad_loss.backward()
        self.opt.step()

5. Run Iterations

(1) Agent & Metric 정의

Agent ( REINFORCE 알고리즘 사용 )

• policy 모델로 MLP (hidden unit=128) 사용

net = MLP(s_dim, a_dim, [128])
agent = REINFORCE(net)

Metric : EMA ( Exponentially Moving Average )

ema = EMAMeter()

(2) Run

• 총 epsiode 횟수 : 10,000회 ( 500회 간격으로 결과 출력 )

n_episode = 10000
print_log = 500

for ep in range(n_episode):
	# (1) Environment 초기화
    s = env.reset()
    
    # (2) epsiode 1회의 cumulative reward
    cum_r = 0
    
    # (3) Terminal state 도착할 때 break
    states = []
    actions = []
    rewards = []
    while True:
        s = to_tensor(s, size=(1, 4))
        a = agent.get_action(s)
        s_, r, done, info = env.step(a.item())
        states.append(s)
        actions.append(a)
        rewards.append(r)

        s = s_
        cum_r += r
        if done:
            break
    ema.update(cum_r)
    
    # (4) (print_log회마다) log 출력 
    if ep % print_log == 0:
        print("Episode {} || EMA: {} ".format(ep, ema.s))

    S = torch.cat(states, dim=0)  # [num. steps x state dim]
    A = torch.stack(actions).squeeze()  # [num. steps]
    R = torch.tensor(rewards)  # torch.tensor [num. steps]
    episode = (S, A, R)
    agent.update_episode(episode, use_norm=True)