[강화학습] REINFORCE 알고리즘 : 코드 구현

 

이번 포스팅에선 REINFORCE 알고리즘을 Pytorch로 간단하게 구현을 해보고자 한다.

다른 RL Algorithms과는 달리 아주 간단하게 예제 표현이 가능하여 어렵지 않다.

본 포스팅을 보기 전, REINFORCE 개념은 확실하게 인지하고 있어야 함을 알린다.

 

 

Last posting
내용: 강화학습 알고리즘인 REINFORCE의 개념과 수식에 대해 살펴보았다.
요약: reward의 합인 목적함수를 최대화하는 최적의 policy 찾는 것이 목표인 알고리즘

 

https://mengu.tistory.com/136

[강화학습] REINFORCE 알고리즘 : 개념 및 수식

 

 

 

1. Problem to solve : CartPole

 

출처: OpenAI

 

설명: CartPole 게임은 하단 박스를 움직여서 막대기가 쓰러지지 않도록 하는 게임이다.
보상: 쓰러지지 않고 버티는 시간만큼 보상을 얻게 된다. 혹은 바닥과 막대기의 위쪽과의 거리가 멀수록 보상을 주는 체계를 구상할 수 있다.

 

 

 

2. 전체 코드

코드는 책 Foundation of Deep Reinforcement Learning(laura.G)를 참고하였다.

from torch.distributions import Categorical
import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

gamma = 0.99

class Pi(nn.Module):
  def __init__(self, in_dim, out_dim):
    super(Pi, self).__init__()
    layers = [
        nn.Linear(in_dim, 64),
        nn.ReLU(),
        nn.Linear(64, out_dim),
        ]
    self.model = nn.Sequential(*layers)
    self.onpolicy_reset()
    self.train() # set training mode
  
  def onpolicy_reset(self):
    self.log_probs = []
    self.rewards = []
  
  def forward(self, x):
    pdparam = self.model(x)
    return pdparam
    
  def act(self, state):
    x = torch.from_numpy(state.astype(np.float32)) # to tensor
    pdparam = self.forward(x) # forward pass
    pd = Categorical(logits=pdparam) # probability distribution
    action = pd.sample() # pi(a|s) in action via pd
    log_prob = pd.log_prob(action) # log_prob of pi(a|s)
    self.log_probs.append(log_prob) # store for training
    return action.item()
  
def train(pi, optimizer):
  # Inner gradient-ascent loop of REINFORCE algorithm
  T = len(pi.rewards)
  rets = np.empty(T, dtype=np.float32) # the returns
  future_ret = 0.0
  # compute the returns efficiently
  for t in reversed(range(T)):
    future_ret = pi.rewards[t] + gamma * future_ret
    rets[t] = future_ret

  rets = torch.tensor(rets)
  log_probs = torch.stack(pi.log_probs)
  loss = - log_probs * rets # gradient term; Negative for maximizing
  loss = torch.sum(loss)
  optimizer.zero_grad()
  loss.backward() # backpropagate, compute gradients
  optimizer.step() # gradient-ascent, update the weights
  return loss

def main():
  env = gym.make('CartPole-v0')
  in_dim = env.observation_space.shape[0] # 4
  out_dim = env.action_space.n # 2
  pi = Pi(in_dim, out_dim) # policy pi_theta for REINFORCE
  optimizer = optim.Adam(pi.parameters(), lr=0.01)

  for epi in range(300):
    state = env.reset()
    for t in range(200): # cartpole max timestep is 200
      action = pi.act(state)
      state, reward, done, _ = env.step(action)
      pi.rewards.append(reward)
      env.render()
      if done:
        break
    loss = train(pi, optimizer) # train per episode
    total_reward = sum(pi.rewards)
    solved = total_reward > 195.0
    pi.onpolicy_reset() # onpolicy: clear memory after training
    
    print(f'Episode {epi}, loss: {loss}, \
    total_reward: {total_reward}, solved: {solved}')
    
if __name__ == '__main__':
  main()

 

약 80줄 정도의 코드로, 지금부터 함수 단위로 쪼개어서 설명하겠다.

위의 코드도 다음과 같은 Algorithm을 따르며, 최적의 ((\pi)) 파라미터를 구하는 것임을 잊지 않는다면 쉽게 이해할 수 있을 것이다.

 

 

 

 

 

2-1. 코드 설명: Class Pi(nn.Module)

Class Pi는 ((\pi_{\theta}))를 코드로 나타낸 것이다. 즉, 학습시킬 인공신경망을 구축하는 코드다.

Pytorch의 nn.Module을 받은 후, __init__() 함수에서 기본적인 세팅을 해주었다.

layer은 [64개의 units으로 구성된 선형신경망-ReLU-선형신경망]이다. 

class Pi(nn.Module):
  def __init__(self, in_dim, out_dim):
    super(Pi, self).__init__()
    layers = [
        nn.Linear(in_dim, 64),
        nn.ReLU(),
        nn.Linear(64, out_dim),
        ]
    self.model = nn.Sequential(*layers)
    self.onpolicy_reset()
    self.train() # set training mode
  
  def onpolicy_reset(self):
    self.log_probs = []
    self.rewards = []
  
  def forward(self, x):
    pdparam = self.model(x)
    return pdparam
    
  def act(self, state):
    x = torch.from_numpy(state.astype(np.float32)) # to tensor
    pdparam = self.forward(x) # forward pass
    pd = Categorical(logits=pdparam) # probability distribution
    action = pd.sample() # pi(a|s) in action via pd
    log_prob = pd.log_prob(action) # log_prob of pi(a|s)
    self.log_probs.append(log_prob) # store for training
    return action.item()

그 밑으로 함수 onpolicy_reset은 training 후에 다시 log_probs, rewards를 초기화하며,

함수 forward는 Agent의 state를 바탕으로 순전파 계산을 수행한다.

함수 act는 state를 인자로 받으며  forward()를 불러와 계산 후, action을 sample로 뽑는다. 또한 action에 따른 ((\bigtriangledown_{theta}))log((\pi_{\theta}(a_{t}|s_{t}))) 도 계산하여 저장해 둔다.

* log_prob = ((\bigtriangledown_{theta}))log((\pi_{\theta}(a_{t}|s_{t})))

 

 

 

 

2-2. 코드 설명: train(pi, optimizer)

함수 train()은 ((\pi))의 파라미터인 ((\theta))를 업데이트하기 위한 loss를 계산한다.
((\pi_{\theta}))에 계산되어 있는 reward를 가져와서 gamma를 통해 discounted reward들을 구한다.

def train(pi, optimizer):
  # Inner gradient-ascent loop of REINFORCE algorithm
  T = len(pi.rewards)
  rets = np.empty(T, dtype=np.float32) # the returns
  future_ret = 0.0
  # compute the returns efficiently
  for t in reversed(range(T)):
    future_ret = pi.rewards[t] + gamma * future_ret
    rets[t] = future_ret

  rets = torch.tensor(rets)
  log_probs = torch.stack(pi.log_probs)
  loss = - log_probs * rets # gradient term; Negative for maximizing
  loss = torch.sum(loss)
  optimizer.zero_grad()
  loss.backward() # backpropagate, compute gradients
  optimizer.step() # gradient-ascent, update the weights
  return loss

그다음, reward와 log_prob를 곱하고, negative 하여 loss sum을 계산한다. 

계산된 loss를 바탕으로 loss.backward()를 계산하고, opimizer를 불러와 최적화를 진행한다.

 

 

 

 

 

2-3. 코드 설명: main()

함수 main은 학습 과정 자체를 나타낸 것이다.

처음 gym 라이브러리를 통해 'CartPole-v0'을 수행할 환경을 불러온다. 그 후 ((\pi)) 신경망과 optimizer 세팅을 해준다.

* in_dim, out_dim은 각각 들어갈 때의 정보 차원(4), 나올 때의 정보 차원(2)을 뜻한다.

 

def main():
  env = gym.make('CartPole-v0')
  in_dim = env.observation_space.shape[0] # 4
  out_dim = env.action_space.n # 2
  pi = Pi(in_dim, out_dim) # policy pi_theta for REINFORCE
  optimizer = optim.Adam(pi.parameters(), lr=0.01)

  for epi in range(300):
    state = env.reset()
    for t in range(200): # cartpole max timestep is 200
      action = pi.act(state)
      state, reward, done, _ = env.step(action)
      pi.rewards.append(reward)
      env.render()
      if done:
        break
    loss = train(pi, optimizer) # train per episode
    total_reward = sum(pi.rewards)
    solved = total_reward > 195.0
    pi.onpolicy_reset() # onpolicy: clear memory after training
    
    print(f'Episode {epi}, loss: {loss}, \
    total_reward: {total_reward}, solved: {solved}')

episode(=epoch)는 300번이다. 매 episode마다 state는 초기화되고, timeste 200번을 돌 때까지 실패하지 않았다면 성공으로 간주한다. 매 timestep마다 ((\pi)) 신경망을 통해 action을 sample 하며, 이에 따른 다음 state, reward를 구하여 저장해 둔다.

timestep이 종료되면, 함수 train()으로 loss를 구하고 업데이트한다. 한편, ((\pi))에서 rewards를 가져와 sum 하여 최종 성공 여부를 판단한다. 그 후, 다시 episode를 시작할 땐 모든 log_prob, reward*를 초기화해 준다.

* 이는 REINFORCE 알고리즘이 on-policy를 따르기 때문이다. 오로지 최근 policy에 의해서만 update 되어야 한다. 예를 들면, 다음 episode 105의 업데이트는 episode 104의 policy에서 비롯되어야 한다.

 

 

 

 

 

3. 결과

 

최초로 성공한 Episode는 98번째였다. 하지만 우연히 성공한 것이기에, 이후부터는 다시 실패하였다. 하지만 이런 성공 경험이 reward 향상에 큰 기여를 했을 것이다.

 

 

 

총 300번의 Episode를 진행했고, total_reward는 점차 200에 수렴했다.

 

 

 

지금까지 REINFORCE 강화학습 알고리즘의 개념과 코드를 살펴보았다. 다음 포스팅에선 다른 알고리즘을 다루도록 하겠다.