def load_checkpoint(file_dir, i_epoch, layer_sizes, input_size, device='cuda'):
checkpoint = torch.load(os.path.join(file_dir, "ckpt_eps%d.pt" % i_epoch),
map_location=device)
policy_net = PolicyNet(layer_sizes).to(device)
value_net = ValueNet(input_size).to(device)
policy_net.load_state_dict(checkpoint["policy_net"])
policy_net.train()
value_net.load_state_dict(checkpoint["value_net"])
value_net.train()
policy_lr = checkpoint["policy_lr"]
valuenet_lr = checkpoint["valuenet_lr"]
policynet_optim = optim.Adam(policy_net.parameters(), lr=policy_lr)
policynet_optim.load_state_dict(checkpoint["policynet_optim"])
valuenet_optim = optim.Adam(value_net.parameters(), lr=valuenet_lr)
valuenet_optim.load_state_dict(checkpoint["valuenet_optim"])
checkpoint.pop("policy_net")
checkpoint.pop("value_net")
checkpoint.pop("policynet_optim")
checkpoint.pop("valuenet_optim")
checkpoint.pop("i_epoch")
checkpoint.pop("policy_lr")
checkpoint.pop("valuenet_lr")
return policy_net, value_net, policynet_optim, valuenet_optim, checkpoint
# observation = env.reset()
# print(observation)
# print(env.observation_space)
MAXSTEP = 100
BATCHSIZE = 16
EPOCH = 1000
GAMMA = 0.99
policy_net = PolicyNet()
value_net = ValueNet()
policy_net.cuda()
value_net.cuda()
opt1 = optim.Adam(policy_net.parameters(), lr=1e-3)
opt2 = optim.Adam(value_net.parameters(), lr=1e-3)
# train one epoch
def train_step():
observ_batch = []
reward_batch = []
action_batch = []
mask_batch = []
policy_net.cpu()
value_net.cpu()
for _ in range(BATCHSIZE):
observ = []
reward = []
class A3CGlobal:
def __init__(self, config):
self.config = config
# 정책신경망 생성
self.actor = PolicyNet(self.config.n_state, self.config.n_action)
self.actor.to(device)
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
lr=self.config.actor_lr)
# 가치신경망 생성
self.critic = ValueNet(self.config.n_state, 1)
self.critic.to(device)
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
lr=self.config.critic_lr)
# 리턴값 계산
def get_returns(self, rewards, done, next_value):
returns = torch.zeros(len(rewards),
dtype=torch.float).to(self.config.device)
R = 0 if done else next_value
for i in reversed(range(0, len(rewards))):
R = rewards[i] + self.config.discount_factor * R
returns[i] = R
return returns
# 각 타임스텝마다 정책신경망과 가치신경망을 업데이트
def train_model(self, states, actions, rewards, next_states, done):
states = torch.tensor(states, dtype=torch.float).to(self.config.device)
actions = torch.tensor(actions,
dtype=torch.float).to(self.config.device)
next_states = torch.tensor(next_states,
dtype=torch.float).to(self.config.device)
next_values = self.critic(next_states).view(-1)
# 리턴값 계산
returns = self.get_returns(rewards, done, next_values[-1])
values = self.critic(states).view(-1)
# 가치신경망 학습
critic_loss = self.train_critic(values, returns)
# 정책신경망 학습
actor_loss = self.train_actor(states, actions, returns - values)
return actor_loss, critic_loss
# 정책신경망을 업데이트하는 함수
def train_actor(self, states, actions, advantages):
policy = self.actor(states)
action_prob = torch.sum(actions * policy, dim=1)
cross_entropy = torch.log(action_prob + 1.e-7) * advantages.detach()
actor_loss = -torch.mean(cross_entropy)
self.actor_optimizer.zero_grad()
actor_loss.backward()
self.actor_optimizer.step()
return actor_loss.item()
# 가치신경망을 업데이트하는 states
def train_critic(self, values, targets):
critic_loss = torch.mean(torch.pow(targets - values, 2))
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
return critic_loss.item()
# GPU 메모리 반납
def close(self):
del self.actor
del self.critic
class A2CAgent:
def __init__(self, config):
self.config = config
# replay memory
self.replay_memory = deque(maxlen=self.config.n_replay_memory)
# 정책신경망 생성
self.actor = PolicyNet(self.config.n_state, self.config.n_action)
self.actor.to(device)
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=self.config.actor_lr)
# 가치신경망 생성
self.critic = ValueNet(self.config.n_state, 1)
self.critic.to(device)
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=self.config.critic_lr)
# 정책신경망의 출력을 받아 확률적으로 행동을 선택
def get_action(self, state):
state = torch.tensor(state, dtype=torch.float).to(device)
policy = self.actor(state)
policy = policy.detach().cpu().numpy()[0]
return np.random.choice(self.config.n_action, 1, p=policy)[0]
# 히스토리 추가
def append_replay(self, state, action, reward, next_state):
act = np.zeros(self.config.n_action)
act[action] = 1
self.replay_memory.append((state, act, reward, next_state))
# 리턴값 계산
def get_returns(self, rewards, done, next_value):
returns = torch.zeros(len(rewards), dtype=torch.float).to(self.config.device)
R = 0 if done else next_value
for i in reversed(range(0, len(rewards))):
R = rewards[i] + self.config.discount_factor * R
returns[i] = R
return returns
# 각 타임스텝마다 정책신경망과 가치신경망을 업데이트
def train_model(self, done):
# 히스토리를 배열 형태로 정렬
replay_memory = np.array(self.replay_memory)
self.replay_memory.clear()
states = np.vstack(replay_memory[:, 0])
actions = list(replay_memory[:, 1])
rewards = list(replay_memory[:, 2])
next_states = list(replay_memory[:, 3])
states = torch.tensor(states, dtype=torch.float).to(self.config.device)
actions = torch.tensor(actions, dtype=torch.float).to(self.config.device)
next_states = torch.tensor(next_states, dtype=torch.float).to(self.config.device)
next_values = self.critic(next_states).view(-1)
# 리턴값 계산
returns = self.get_returns(rewards, done, next_values[-1])
values = self.critic(states).view(-1)
# 가치신경망 학습
critic_loss = self.train_critic(values, returns)
# 정책신경망 학습
actor_loss = self.train_actor(states, actions, returns - values)
return actor_loss, critic_loss
# 정책신경망을 업데이트하는 함수
def train_actor(self, states, actions, advantages):
policy = self.actor(states)
action_prob = torch.sum(actions * policy, dim=1)
cross_entropy = torch.log(action_prob + 1.e-7) * advantages.detach()
actor_loss = -torch.mean(cross_entropy)
self.actor_optimizer.zero_grad()
actor_loss.backward()
self.actor_optimizer.step()
return actor_loss.item()
# 가치신경망을 업데이트하는 states
def train_critic(self, values, targets):
critic_loss = torch.mean(torch.pow(targets - values, 2))
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
return critic_loss.item()
# model의 weight를 파일로 저장
def save(self):
torch.save(self.actor.state_dict(), self.config.save_file + ".actor")
torch.save(self.critic.state_dict(), self.config.save_file + ".critic")
# 파일로 부터 model의 weight를 읽어 옴
def load(self):
self.actor.load_state_dict(torch.load(self.config.save_file + ".actor"))
self.critic.load_state_dict(torch.load(self.config.save_file + ".critic"))
# GPU 메모리 반납
def close(self):
del self.actor
del self.critic
env = env.unwrapped
# Get device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Current usable device is: ", device)
# Create the model
policy_net = PolicyNet(layer_sizes, action_lim).to(device) # Policy network
value_net = ValueNet(input_size).to(device) # Value network
# Set up memory
memory = Memory(capacity, device)
# Set up optimizer
policynet_optimizer = optim.Adam(policy_net.parameters(), lr=policy_lr)
valuenet_optimizer = optim.Adam(value_net.parameters(), lr=valuenet_lr)
###################################################################
# Start training
# Dictionary for extra training information to save to checkpoints
training_info = {
"epoch mean durations": [],
"epoch mean rewards": [],
"max reward achieved": 0,
"past %d epochs mean reward" % num_avg_epoch: 0,
"value net loss": []
}
# Batch that records trajectories
class DQNAgent:
def __init__(self, config):
self.config = config
self.epsilon = config.epsilon
# replay memory
self.replay_memory = deque(maxlen=self.config.n_replay_memory)
# 가치신경망 생성
self.model = ValueNet(self.config.n_state, self.config.n_action)
self.model.to(device)
self.model_optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.config.learning_rate)
# 정책신경망의 출력을 받아 확률적으로 행동을 선택
def get_action(self, state):
if np.random.rand() <= self.epsilon:
return random.randrange(self.config.n_action)
else:
state = torch.tensor(state,
dtype=torch.float).to(self.config.device)
output = self.model(state)
return output.argmax().item()
# 히스토리 추가
def append_replay(self, state, action, reward, next_state, done):
self.replay_memory.append((state, action, reward, next_state, done))
# 각 타임스텝마다 정책신경망과 가치신경망을 업데이트
def train_model(self):
# 학습이 계속 될 수 록 탐험 학률을 줄여 줌
if self.epsilon > self.config.epsilon_min:
self.epsilon *= self.config.epsilon_decay
# 히스토리를 배열 형태로 정렬
replay_memory = np.array(
random.sample(self.replay_memory, self.config.n_batch))
states = np.vstack(replay_memory[:, 0])
actions = list(replay_memory[:, 1])
rewards = list(replay_memory[:, 2])
next_states = list(replay_memory[:, 3])
dones = list(replay_memory[:, 4])
states = torch.tensor(states, dtype=torch.float).to(device)
next_states = torch.tensor(next_states, dtype=torch.float).to(device)
targets = self.model(states)
next_values = self.model(next_states)
for i in range(len(targets)):
if dones[i]:
targets[i][actions[i]] = rewards[i] # Vt = Rt+1
else:
targets[i][
actions[i]] = rewards[i] + self.config.discount_factor * (
torch.max(next_values[i])) # Vt = Rt+1 + rVt+1
loss = self.train_value(states, targets)
return loss
# 가치신경망을 업데이트하는 함수
def train_value(self, states, targets):
values = self.model(states)
loss = torch.mean(torch.pow(targets - values, 2))
self.model_optimizer.zero_grad()
loss.backward()
self.model_optimizer.step()
return loss.item()
# model의 weight를 파일로 저장
def save(self):
torch.save(self.model.state_dict(), self.config.save_file)
# 파일로 부터 model의 weight를 읽어 옴
def load(self):
self.model.load_state_dict(torch.load(self.config.save_file))
# GPU 메모리 반납
def close(self):
del self.model
class SAC:
def __init__(self, env, gamma, tau, buffer_maxlen, value_lr, q_lr, policy_lr):
self.env = env
self.state_dim = env.observation_space.shape[0]
self.action_dim = env.action_space.shape[0]
self.action_range = [env.action_space.low, env.action_space.high]
# hyperparameters
self.gamma = gamma
self.tau = tau
# initialize networks
self.value_net = ValueNet(self.state_dim).to(device)
self.target_value_net = ValueNet(self.state_dim).to(device)
self.q1_net = SoftQNet(self.state_dim, self.action_dim).to(device)
self.q2_net = SoftQNet(self.state_dim, self.action_dim).to(device)
self.policy_net = PolicyNet(self.state_dim, self.action_dim).to(device)
# Load the target value network parameters
for target_param, param in zip(self.target_value_net.parameters(), self.value_net.parameters()):
target_param.data.copy_(self.tau * param + (1 - self.tau) * target_param)
# Initialize the optimizer
self.value_optimizer = optim.Adam(self.value_net.parameters(), lr=value_lr)
self.q1_optimizer = optim.Adam(self.q1_net.parameters(), lr=q_lr)
self.q2_optimizer = optim.Adam(self.q2_net.parameters(), lr=q_lr)
self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=policy_lr)
# Initialize thebuffer
self.buffer = ReplayBeffer(buffer_maxlen)
def get_action(self, state):
action = self.policy_net.action(state)
action = action * (self.action_range[1] - self.action_range[0]) / 2.0 + \
(self.action_range[1] + self.action_range[0]) / 2.0
return action
def update(self, batch_size):
state, action, reward, next_state, done = self.buffer.sample(batch_size)
new_action, log_prob = self.policy_net.evaluate(state)
# V value loss
value = self.value_net(state)
new_q1_value = self.q1_net(state, new_action)
new_q2_value = self.q2_net(state, new_action)
next_value = torch.min(new_q1_value, new_q2_value) - log_prob
value_loss = F.mse_loss(value, next_value.detach())
# Soft q loss
q1_value = self.q1_net(state, action)
q2_value = self.q2_net(state, action)
target_value = self.target_value_net(next_state)
target_q_value = reward + done * self.gamma * target_value
q1_value_loss = F.mse_loss(q1_value, target_q_value.detach())
q2_value_loss = F.mse_loss(q2_value, target_q_value.detach())
# Policy loss
policy_loss = (log_prob - torch.min(new_q1_value, new_q2_value)).mean()
# Update v
self.value_optimizer.zero_grad()
value_loss.backward()
self.value_optimizer.step()
# Update Soft q
self.q1_optimizer.zero_grad()
self.q2_optimizer.zero_grad()
q1_value_loss.backward()
q2_value_loss.backward()
self.q1_optimizer.step()
self.q2_optimizer.step()
# Update Policy
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
# Update target networks
for target_param, param in zip(self.target_value_net.parameters(), self.value_net.parameters()):
target_param.data.copy_(self.tau * param + (1 - self.tau) * target_param)