개요
REINFORCE with baseline 알고리즘을 PyTorch 코드로 구현해보았으며 Colab에서 코드 실행 및 결과 확인이 가능합니다.
설명
REINFORCE with baseline 알고리즘은 REINFORCE 알고리즘에 baseline function 적용을 통해 학습 안정성을 개선한 알고리즘입니다. 기존 REINFORCE 알고리즘은 학습이 잘 되다가 다시 성능이 낮아지는 경향이 있지만 baseline이 적용된 경우 비교적 안정적인 모습을 보입니다.
Baseline function으로 state-value function역할을 하는 ValueNetwork를 학습시켜 사용하였습니다.
def update_weights_with_baseline(self):
G_t = 0
self.P_optimizer.zero_grad()
self.V_optimizer.zero_grad()
for gradient_policy_a_s, r, v in self.trajectory[::-1]: #cumulative future rewards들의 합 계산을 위해 [::-1]
G_t = 0.99 * G_t + r
delta = G_t - v
V_loss = -1 * delta.item()* v # V_loss = delta**2
V_loss.backward()
P_loss = -1 * delta.item() * gradient_policy_a_s # -1 for gradient ascent
P_loss.backward()
self.P_optimizer.step()
self.V_optimizer.step()
self.trajectory.clear()
return
계산그래프상 PolicyNetwork랑 ValueNetwork의 충돌을 막기 위해서 delta는 .item()으로 값만 사용해 loss를 구성해야 합니다.
확실히 학습 안정성 및 속도 측면에서 우월함을 보입니다. 결과로 돌려보는 게임에서도 움직임이 더 안정적입니다.
REINFORCE with baseline(cartpole)
REINFORCE(cartpole)
Lunalander 게임의 경우에도 episode를 길게 두고 학습시키면 좋은 성능으로 수렴하는 모습을 보입니다.
REINFORCE with baseline(lunarlander-v2)
Full code
import gym
import torch
import torch.nn as nn
from torch.distributions import Categorical
import matplotlib.pyplot as plt
!pip install gym[classic_control]
#pip install gym[box2d] #for lunarlander
!apt update
!apt install xvfb
!pip install pyvirtualdisplay
!pip install gym-notebook-wrapper
import gnwrapper
!nvidia-smi
print(torch.cuda.is_available())
class Agent(nn.Module):
def __init__(self, input_dim, output_dim, width):
super().__init__()
self.P = PolicyNetwork(input_dim, output_dim, width)
self.P.to(device)
self.P.train()
self.P_optimizer = torch.optim.Adam(self.P.parameters(), lr=0.0003)
self.V = ValueNetwork(input_dim, output_dim, width)
self.V.to(device)
self.V.train()
self.V_optimizer = torch.optim.Adam(self.V.parameters(), lr=0.01)
self.trajectory = []
self.env = gym.make(game_name)
def self_play(self, max_timestep=1000000):
game_score = 0
state = self.env.reset() # env 시작
for _ in range(max_timestep):
output = self.P(torch.from_numpy(state).float().to(device)) # inference
inferenced_v = self.V(torch.from_numpy(state).float().to(device))
prob_distribution = Categorical(output) #확률분포 표현
action = prob_distribution.sample() #확률분포로부터 action 선택
state, r, done, _ = self.env.step(action.item()) # env 진행
self.trajectory.append((prob_distribution.log_prob(action), r, inferenced_v)) # a,r,v 저장
game_score += r
if done:
break
return game_score
def update_weights_with_baseline(self):
G_t = 0
self.P_optimizer.zero_grad()
self.V_optimizer.zero_grad()
for gradient_policy_a_s, r, v in self.trajectory[::-1]: #cumulative future rewards들의 합 계산을 위해 [::-1]
G_t = 0.99 * G_t + r
delta = G_t - v
V_loss = -1 * delta.item()* v # V_loss = delta**2
V_loss.backward()
P_loss = -1 * delta.item() * gradient_policy_a_s # -1 for gradient ascent
P_loss.backward()
self.P_optimizer.step()
self.V_optimizer.step()
self.trajectory.clear()
return
class PolicyNetwork(nn.Module):
def __init__(self, input_dim, output_dim, width):
super().__init__()
self.layer1 = torch.nn.Linear(input_dim, width)
self.layer2 = torch.nn.Linear(width, width)
self.layer3 = torch.nn.Linear(width, output_dim)
def forward(self, x):
x = self.layer1(x)
x = torch.nn.functional.relu(x)
x = self.layer2(x)
x = torch.nn.functional.relu(x)
x = self.layer3(x)
x = torch.nn.functional.softmax(x, dim=0)
return x
class ValueNetwork(nn.Module):
def __init__(self, input_dim, output_dim, width):
super().__init__()
self.layer1 = torch.nn.Linear(input_dim, width)
self.layer2 = torch.nn.Linear(width, width)
self.layer3 = torch.nn.Linear(width, 1)
def forward(self, x):
x = self.layer1(x)
x = torch.nn.functional.relu(x)
x = self.layer2(x)
x = torch.nn.functional.relu(x)
x = self.layer3(x)
return x
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
score_arr = []
game_name = 'CartPole-v1' #LunarLander-v2
env = gym.make(game_name)
agent = Agent(env.observation_space.shape[0], env.action_space.n, 128)
print(agent)
env.close()
#Self play 및 weight update
episode_nums = 500 #LunarLander-v2는 1500
for i in range(episode_nums):
game_score = agent.self_play() # play길이 제한 필요시 (max_timestep=값)
score_arr.append(game_score)
agent.update_weights_with_baseline()
if i%50==0 : print('episode', i)
torch.save(agent.state_dict(), 'weights.pt')
agent.env.close()
#Episode별 얻은 score
plt.plot(score_arr, label ='score')
plt.legend(loc='upper left')
#학습된 모델로 게임 play한 영상
agent.load_state_dict(torch.load("weights.pt"))
env = gnwrapper.LoopAnimation(gym.make(game_name))
state = env.reset()
for _ in range(200):
with torch.no_grad():
output = agent.P(torch.from_numpy(state).float().to(device)) # inference
prob_distribution = Categorical(output) #확률분포 표현
action = prob_distribution.sample() #확률분포로부터 action 선택
env.render()
state, rew, done, _ = env.step(action.item())
if done:
state = env.reset()
env.display()
댓글
댓글 쓰기