아티클 정리 - Policy Gradient Reinforcement Learning in PyTorch

안녕하세요. 오늘은 꽤나 흥미로운 포스팅을 하나 준비해왔습니다. 최근에 딥 러닝의 추세라고 볼 수 있는 강화학습과 관련된 포스팅입니다. 본문은 아래와 같습니다. 링크

Policy Gradient Reinforcement Learning in PyTorch

강화학습에 대해서 설명하기 위해서 먼저 environment와 agent에 대해서 설명해야합니다. agent는 딥 러닝으로 생각해보면 모델이 됩니다. agent는 항상 특정 state에서 최상의 reward를 얻기 위한 action을 취하게 됩니다. 그 reward와 다음 state에 대한 정보는 environment와 상호작용하면서 알게 되는 것이죠. 이때, action은 항상 policy로부터 오게 됩니다. 바로 오늘 포스팅의 핵심 단어겠죠.

 

이러한 강화학습은 기본적으로 greedy algorithm에 그 뿌리를 두고 있습니다. 예를 들어서 현재 state에서 {"왼쪽", "오른쪽"}이라는 action 중 한 가지를 취할 수 있습니다. 그러면 agent는 "왼쪽" action을 선택해야 바로 받을 수 있는 reward가 더 크다고 가정하면 agent는 "왼쪽"이라는 action을 선택하느 것이죠. 하지만!! 강화학습은 이보다 좀 더 general한 것을 원합니다. 바로 현재의 reward가 아닌 앞으로 받을 모든 reward의 평균인 기댓값을 최대화하는 방향으로 action을 선택하려는 것이죠. 기존의 greedy algorithm은 간단하지만 항상 최적성을 보장하지 않는다는 단점이 있습니다. 하지만, environment가 finite Markov Decision Process일 경우에는 Bellman optimality equation에 따라서 다음 state를 위한 action을 예측하여 다음 reward의 기댓값이 최상이 되는 optimal solution이 유일하게 존재할 수 있으며 찾을 수 있다는 것입니다. 

 

또한 강화학습은 최상의 reward를 얻기 위해서 몇 가지 방식을 구현할 수 있습니다. 그 중에 하나가 Policy gradient라는 것만 알아두시면 될 거 같습니다.

 

보시기 전에 참고하시면 이 포스팅은 2018년도 Pytorch 버전을 사용하고 있기 때문에 현재의 Pytorch와는 다른 점이 있을 수 있습니다. 그렇기 때문에 본문의 코드를 그대로 실행하고자 하시는 분들은 pytorch 0.3.1 버전을 사용하시길 바랍니다. 저는 현재 pytorch 버전의 코드를 사용하기 위해서 코드를 조금 뜯어고쳤는 데 제대로 학습되지 않았으니 마음이 아픕니다 ㅠㅠ

1. CartPole

파이썬에서는 강화학습은 environment를 지원하기 위한 라이브러리인 gym을 지원하고 있습니다. 이 gym이라는 라이브러리는 일전에 일론 머스크가 설립한 인공지능 연구 회사인 openAI의 결과입니다(지금은 일론 머스크가 없다는 점...). gym에서는 다양한 environment를 지원하고 있습니다. 그 중에서 오늘 다룰 environment는 "CartPole-v1"입니다. 어떤 느낌인지는 아래의 사진을 보시면 될 거 같습니다.

보시면 검은색 상자(Cart)가 막대기(Pole)의 균형을 아주 잘 유지하고 있습니다. 바로 이것이 목표입니다. 검은색 상자는 왼쪽 또는 오른쪽으로만 움직이면서 막대기의 균형을 최대한 오랫동안 유지하려는 것입니다. 아주 간단한 문제 같지만 컴퓨터에게 시키려고 하니 꽤나 갑갑합니다. 그래서 오늘 강화학습을 이용해서 이 균형을 최대한 유지해보도록 하겠습니다!

 

2. Pytorch

오늘 사용할 딥 러닝 프레임워크는 Pytorch입니다. Pytorch는 Tensorflow 1.0 버전에 비해서 진입장벽이 낮아서 초보자 분들이 배우기에 적합한 프레임워크 중에 하나입니다. 물론 현재는 tensorflow 2.0 버전이 나와서 tensorflow 역시 기존의 1.0 버전보다는 많이 쉬워졌지만 그럼에도 불구하고 Tensorflow에 비해서 자유도가 높기 때문에 딥 러닝 프레임워크의 양대산맥이라고 할 수 있습니다.

 

3. Policy Gradients

오늘의 핵심입니다!! Policy Gradient는 수많은 강화학습의 알고리즘 중에 하나일 뿐입니다. 바로 일전에서 강화학습에 대해서 간단히 설명드릴 때 Policy에 대해서 기억하실겁니다. agent는 현재 state부터 마지막 state(terminal state)까지의 reward의 기댓값을 최대화하기위한 action을 선택한다고 하였습니다. 여기서 알고리즘이 2개로 나뉘게 됩니다. Policy가 고정된 경우와 Policy를 바꾸는 경우입니다. Policy가 고정된 경우에는 agent는 한정된 action을 취할 수 밖에 없지만 Policy를 바꾸게 되면 agent는 좀 더 많은 종류의 action을 선택할 수 있을 것입니다. 

 

오늘 포스팅을 이해하기 위해서는 아래의 식을 먼저 이해해야합니다.

$\phi$는 Policy function이라고 부릅니다. function에 대해서 생각해보면 function은 항상 어떤 input을 받으면 output을 얻는 것입니다. 이 경우에는 s, a라는 것을 input으로 가지는데요. s는 state, a는 action을 의미합니다. 즉, 특정 state와 특정 action을 입력으로 받아서 어떤 값을 출력하게 됩니다. 이 출력값은 바로 다음 state에서 취할 action입니다. 마지막으로 $\theta$는 저희가 딥 러닝을 사용해서 구현할텐데 딥 러닝 모델의 파라미터라고 보시면 될 거 같습니다.   

 

정리하면 딥 러닝 모델을 이용해서 특정 state, 특정 action을 입력받아 다음에 취할 새로운 action을 취하는 함수라고 어렴풋이 이해하시면 될 거 같습니다! 간단하죠?

 

4. 모델 구성

모델은 그리 복잡하지 않습니다. 128개의 뉴런을 가지는 hidden layer 1개와 보존율 0.6의 Dropout로만 구성되어 있는 단순한 feed-forward network를 사용합니다. 이때, optimizer는 Adam, learning rate는 0.01입니다.

 

솔직히 말해서 강화학습에 dropout을 적용하게 되면 이후의 policy의 변화에 어떤 변화를 줄 지는 저로써는 잘 모르겠습니다. 이는 이후에 논문을 읽게 되면 정리하도록 하겠습니다. 하지만 이번 포스팅의 설명을 보면 dropout을 적용하게 되면 그 성능이 크게 향상된다고 합니다!! 이제, 이 코드를 실행하기 전에 반드시 pip install gym을 이용해서 gym 라이브러리를 설치해주셔야합니다. 

 

env = gym.make('CartPole-v1')
env.seed(1); torch.manual_seed(1);


#Hyperparameters
learning_rate = 0.01
gamma = 0.99


class Policy(nn.Module):
    def __init__(self):
        super(Policy, self).__init__()
        self.state_space = env.observation_space.shape[0]
        self.action_space = env.action_space.n

        self.l1 = nn.Linear(self.state_space, 128, bias=False)
        self.l2 = nn.Linear(128, self.action_space, bias=False)

        self.gamma = gamma

        # Episode policy and reward history
        self.policy_history = Variable(torch.Tensor([]), requires_grad=True)
        self.reward_episode = []
        # Overall reward and loss history
        self.reward_history = []
        self.loss_history = []

    def forward(self, x):
        model = torch.nn.Sequential(
            self.l1,
            nn.Dropout(p=0.6),
            nn.ReLU(),
            self.l2,
            nn.Softmax(dim=-1)
        )
        return model(x)

policy = Policy()
optimizer = optim.Adam(policy.parameters(), lr=learning_rate)

 

이 코드를 부분으로 짤라서 설명을 하도록 하겠습니다.

import gym
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm, trange
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.distributions import Categorical

먼저, 필요한 라이브러리들을 import 합니다.

env = gym.make('CartPole-v1')
env.seed(1); torch.manual_seed(1);

바로 이전에 설명드린 것처럼 openAI의 gym 라이브러리에서는 많은 environment 실습 환경을 제공해주고 있습니다(Thank you openAI!). 오늘은 그 중에서 "CartPole-v1"을 이용하도록 하겠습니다.

# Hyperparameters
learning_rate = 0.01
gamma = 0.99

그리고 하이퍼파라미터를 정의합니다. 강화학습에서는 gamma라는 하이퍼파라미터가 존재하는 데 이 값은 커질수록 미래의 보상을 보고 결정한다는 의미를 같습니다. 그리고 0에 가까워질수록 단순 greedy algorithm과 유사해지는 것이죠. 보통 $k$번째 reward $R_{k}$를 total reward에 더할 때 $\gamma \times \R_{k}$를 더하게 되는 데 이유는 보통 reward는 양수로 이루어져 있기 때문에 무한히 반복한다고 가정하면 발산하게 됩니다. 따라서 $\gamma$를 곱해서 더함으로써 수렴성까지 보장할 수 있는 것이죠. 이를 discount factor update라고 말합니다. 이를 식으로 표현하면 아래와 같습니다.

class Policy(nn.Module) :
    def __init__(self):
        super(Policy, self).__init__()

        self.state_space = env.observation_space.shape[0]
        self.action_space = env.action_space.n

        self.linear1 = nn.Linear(self.state_space, 128, bias=False)
        self.linear2 = nn.Linear(128, self.action_space, bias=False)

        self.gamma = gamma

        # Episode policy and reward history
        self.policy_history = Variable(torch.Tensor())
        self.reward_episode = []

        # Overall reward and loss history
        self.reward_history = []
        self.loss_history = []

    def forward(self, x):
        model = nn.Sequential(
            self.linear1,
            nn.Dropout(0.6),
            nn.ReLU(),
            self.linear2,
            nn.Softmax(dim=-1)
        )

        return model(x)

그 다음은 모델을 정의하게 되는 데 state_space은 CartPole-v1이라는 environment가 가질 수 있는 state의 총 수를 의미합니다. 그리고 action_space는 각 state별로 agent가 취할 수 있는 action의 수입니다.  이를 단순하게 선형으로 연결한 뒤 dropout을 적용하는 것을 볼 수 있습니다. 구성된 모델을 출력해보면 아래와 같은 결과를 얻을 수 있습니다.

policy = Policy()
optimizer = optim.Adam(policy.parameters(), lr=learning_rate)
print(policy)

Policy(
  (linear1): Linear(in_features=4, out_features=128, bias=False)
  (linear2): Linear(in_features=128, out_features=2, bias=False)
)

5. action selection

저희의 environment는 왼쪽, 또는 오른쪽의 action만 취할 수 있기 때문에 단순하게 어떤 방향으로 움직일지에 대한 action만 선택해주면 됩니다. 이때, "왼쪽으로 가라", 또는 "오른쪽으로 가라"라는 것을 확률적으로 결정하게 된다는 것에 유의해주시길 바랍니다.

def select_action(state):
    # Select an action (0 or 1) by running policy model and choosing based on the probabilities in state
    state = torch.from_numpy(state).type(torch.FloatTensor)
    state = policy(Variable(state))
    c = Categorical(state)
    action = c.sample()

    # Add log probability of our chosen action to our history
    if policy.policy_history.dim() != 0:
        policy.policy_history = torch.cat([policy.policy_history, torch.Tensor([c.log_prob(action)])])
    else:
        policy.policy_history = (c.log_prob(action))
    return action

6. update policy

드디어 이번 포스팅의 핵심입니다. 저희는 이제 Monte-Carlo Policy Gradient method를 이용해서 policy를 update하도록 하겠습니다. 이를 위해서는 아래의 방정식을 이해해야합니다.

$\phi_{\theta}(s_{t}, a_{t})$은 이전에 설명드린 것처럼 특정 state $s_{t}$, 특정 action $a_{t}$가 주어졌을 때 다음 state $s_{t+1}$에서 취할 action $a_{t+1}$을 확률적으로 선택하는 함수입니다. $v_{t}$은 $s_{t}$까지 도달했을 때 지금까지 받은 reward의 총합이라고 보면 됩니다. 이 둘을 곱하게 되면 어떻게 될까요? 바로 reward의 기댓값을 얻을 수 있을 것입니다(기댓값의 정의를 천천히 읽어보시면 됩니다.). 이제, 이를 모델의 파라미터인 $\theta$에 대해서 미분을 하고 learning_rate를 곱하여 매개변수를 update하는 것이죠. 코드를 보도록 하겠습니다.

def update_policy():
    R = 0
    rewards = []

    # Discount future rewards back to the present using gamma
    for r in policy.reward_episode[::-1]:
        R = r + policy.gamma * R
        rewards.insert(0, R)

    # Scale rewards
    rewards = torch.FloatTensor(rewards)
    rewards = (rewards - rewards.mean()) / (rewards.std() + np.finfo(np.float32).eps)

    # Calculate loss
    loss = (torch.sum(torch.mul(policy.policy_history, Variable(rewards, requires_grad=True)).mul(-1), -1))

    # Update network weights
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Save and intialize episode history counters
    policy.loss_history.append(loss.item())
    policy.reward_history.append(np.sum(policy.reward_episode))
    policy.policy_history = Variable(torch.Tensor())
    policy.reward_episode = []

7. training

이제는 실제로 학습을 해보도록 하겠습니다. 아래의 코드를 참조해주시길 바랍니다.

def main(episodes):
    running_reward = 10
    for episode in range(episodes):
        state = env.reset()  # Reset environment and record the starting state
        done = False

        for time in range(1000):
            action = select_action(state)
            # Step through environment using chosen action
            state, reward, done, _ = env.step(action.item())

            # Save reward
            policy.reward_episode.append(reward)
            if done:
                break

        # Used to determine when the environment is solved.
        running_reward = (running_reward * 0.99) + (time * 0.01)

        update_policy()

        if episode % 50 == 0:
            print('Episode {}\tLast length: {:5d}\tAverage length: {:.2f}'.format(episode, time, running_reward))

        if running_reward > env.spec.reward_threshold:
            print("Solved! Running reward is now {} and the last episode runs to {} time steps!".format(running_reward,
                                                                                                        time))
            break

8. results

하지만 그 결과를 보면 그렇게 썩 좋지 않습니다. 그 이유는 아무래도 번역한 포스팅이 2년전인 2018년 포스팅이다보니 현재 pytorch의 버전과 상당히 달랐습니다. 그래서 몇 개의 코드를 좀 뜯어고쳤는 데 아무래도 잘못 뜯어고친거 같습니다. 결론은!! 일반적으로 policy gradient 기반 강화학습은 좋은 성능을 낼 수 있다는 것입니다. 마지막으로 해당 포스팅의 결과를 보고 마치도록 하겠습니다. 원래는 아래와 같이 잘 나와야되는데 대체 어디서 잘못짠것인지 모르겠지만 아시는 분들은 댓글을... ㅠㅠ

 

오늘의 교훈...은 포스팅을 볼꺼면 최신 버전의 포스팅을 보자는 것입니다...ㅠㅠ