def __init__(self, net_dim, state_dim, action_dim, learning_rate=1e-4):
super().__init__()
self.target_entropy = np.log(action_dim)
self.alpha_log = torch.tensor(
(-np.log(action_dim) * np.e, ),
dtype=torch.float32,
requires_grad=True,
device=self.device) # trainable parameter
self.act = ActorSAC(net_dim, state_dim, action_dim).to(self.device)
self.act_target = deepcopy(self.act)
self.cri = CriticTwin(
int(net_dim * 1.25),
state_dim,
action_dim,
).to(self.device)
self.cri_target = deepcopy(self.cri)
self.criterion = torch.nn.SmoothL1Loss()
self.optimizer = torch.optim.Adam([{
'params': self.act.parameters(),
'lr': learning_rate * 0.75
}, {
'params': self.cri.parameters(),
'lr': learning_rate * 1.25
}, {
'params': (self.alpha_log, ),
'lr': learning_rate
}])
def init(self, net_dim, state_dim, action_dim, if_per=False):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.target_entropy *= np.log(action_dim)
self.alpha_log = torch.tensor((-np.log(action_dim) * np.e,), dtype=torch.float32,
requires_grad=True, device=self.device) # trainable parameter
self.alpha_optimizer = torch.optim.Adam((self.alpha_log,), self.learning_rate)
self.cri = CriticTwin(int(net_dim * 1.25), state_dim, action_dim, self.if_use_dn).to(self.device)
self.cri_target = deepcopy(self.cri)
self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), self.learning_rate)
self.act = ActorSAC(net_dim, state_dim, action_dim, self.if_use_dn).to(self.device)
self.act_optimizer = torch.optim.Adam(self.act.parameters(), self.learning_rate)
self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')
if if_per:
self.get_obj_critic = self.get_obj_critic_per
else:
self.get_obj_critic = self.get_obj_critic_raw
class AgentModSAC(AgentSAC): # Modified SAC using reliable_lambda and TTUR (Two Time-scale Update Rule)
def __init__(self):
super().__init__()
self.if_use_dn = True
self.obj_c = (-np.log(0.5)) ** 0.5 # for reliable_lambda
def init(self, net_dim, state_dim, action_dim, if_per=False):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.target_entropy *= np.log(action_dim)
self.alpha_log = torch.tensor((-np.log(action_dim) * np.e,), dtype=torch.float32,
requires_grad=True, device=self.device) # trainable parameter
self.alpha_optimizer = torch.optim.Adam((self.alpha_log,), self.learning_rate)
self.cri = CriticTwin(int(net_dim * 1.25), state_dim, action_dim, self.if_use_dn).to(self.device)
self.cri_target = deepcopy(self.cri)
self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), self.learning_rate)
self.act = ActorSAC(net_dim, state_dim, action_dim, self.if_use_dn).to(self.device)
self.act_optimizer = torch.optim.Adam(self.act.parameters(), self.learning_rate)
self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')
if if_per:
self.get_obj_critic = self.get_obj_critic_per
else:
self.get_obj_critic = self.get_obj_critic_raw
def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):
buffer.update_now_len_before_sample()
alpha = self.alpha_log.exp().detach()
update_a = 0
for update_c in range(1, int(buffer.now_len / batch_size * repeat_times)):
'''objective of critic (loss function of critic)'''
obj_critic, state = self.get_obj_critic(buffer, batch_size, alpha)
self.obj_c = 0.995 * self.obj_c + 0.0025 * obj_critic.item() # for reliable_lambda
self.cri_optimizer.zero_grad()
obj_critic.backward()
self.cri_optimizer.step()
self.soft_update(self.cri_target, self.cri, self.soft_update_tau)
'''objective of alpha (temperature parameter automatic adjustment)'''
action_pg, logprob = self.act.get_action_logprob(state) # policy gradient
obj_alpha = (self.alpha_log * (logprob - self.target_entropy).detach()).mean()
self.alpha_optimizer.zero_grad()
obj_alpha.backward()
self.alpha_optimizer.step()
with torch.no_grad():
self.alpha_log[:] = self.alpha_log.clamp(-20, 2)
alpha = self.alpha_log.exp().detach()
'''objective of actor using reliable_lambda and TTUR (Two Time-scales Update Rule)'''
reliable_lambda = np.exp(-self.obj_c ** 2) # for reliable_lambda
if_update_a = (update_a / update_c) < (1 / (2 - reliable_lambda))
if if_update_a: # auto TTUR
update_a += 1
# if reliable_lambda > 0.02:
q_value_pg = torch.min(*self.cri_target.get_q1_q2(state, action_pg))
obj_actor = -(q_value_pg + logprob * alpha.detach()).mean() * reliable_lambda
self.act_optimizer.zero_grad()
obj_actor.backward()
self.act_optimizer.step()
return alpha.item(), self.obj_c
class AgentSAC(AgentBase):
def __init__(self):
super().__init__()
self.target_entropy = None
self.alpha_log = None
self.alpha_optimizer = None
self.target_entropy = 1.0 # * np.log(action_dim)
def init(self, net_dim, state_dim, action_dim, if_per=False):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.target_entropy *= np.log(action_dim)
self.alpha_log = torch.tensor((-np.log(action_dim) * np.e,), dtype=torch.float32,
requires_grad=True, device=self.device) # trainable parameter
self.alpha_optimizer = torch.optim.Adam((self.alpha_log,), self.learning_rate)
self.cri = CriticTwin(net_dim, state_dim, action_dim).to(self.device)
self.cri_target = deepcopy(self.cri)
self.cri_optimizer = torch.optim.Adam(self.cri.parameters(), lr=self.learning_rate)
self.act = ActorSAC(net_dim, state_dim, action_dim).to(self.device)
self.act_optimizer = torch.optim.Adam(self.act.parameters(), lr=self.learning_rate)
self.criterion = torch.nn.SmoothL1Loss(reduction='none' if if_per else 'mean')
if if_per:
self.get_obj_critic = self.get_obj_critic_per
else:
self.get_obj_critic = self.get_obj_critic_raw
def select_action(self, state) -> np.ndarray:
states = torch.as_tensor((state,), dtype=torch.float32, device=self.device).detach_()
action = self.act.get_action(states)[0]
return action.cpu().numpy()
def update_net(self, buffer, target_step, batch_size, repeat_times) -> (float, float):
buffer.update_now_len_before_sample()
alpha = self.alpha_log.exp().detach()
obj_critic = None
for _ in range(int(target_step * repeat_times)):
'''objective of critic'''
obj_critic, state = self.get_obj_critic(buffer, batch_size, alpha)
self.cri_optimizer.zero_grad()
obj_critic.backward()
self.cri_optimizer.step()
self.soft_update(self.cri_target, self.cri, self.soft_update_tau)
'''objective of alpha (temperature parameter automatic adjustment)'''
action_pg, logprob = self.act.get_action_logprob(state) # policy gradient
obj_alpha = (self.alpha_log * (logprob - self.target_entropy).detach()).mean()
self.alpha_optimizer.zero_grad()
obj_alpha.backward()
self.alpha_optimizer.step()
'''objective of actor'''
alpha = self.alpha_log.exp().detach()
obj_actor = -(torch.min(*self.cri_target.get_q1_q2(state, action_pg)) + logprob * alpha).mean()
self.act_optimizer.zero_grad()
obj_actor.backward()
self.act_optimizer.step()
return alpha.item(), obj_critic.item()
def get_obj_critic_raw(self, buffer, batch_size, alpha):
with torch.no_grad():
reward, mask, action, state, next_s = buffer.sample_batch(batch_size)
next_a, next_logprob = self.act.get_action_logprob(next_s)
next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))
q_label = reward + mask * (next_q + next_logprob * alpha)
q1, q2 = self.cri.get_q1_q2(state, action) # twin critics
obj_critic = self.criterion(q1, q_label) + self.criterion(q2, q_label)
return obj_critic, state
def get_obj_critic_per(self, buffer, batch_size, alpha):
with torch.no_grad():
reward, mask, action, state, next_s, is_weights = buffer.sample_batch(batch_size)
next_a, next_logprob = self.act.get_action_logprob(next_s)
next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))
q_label = reward + mask * (next_q + next_logprob * alpha)
q1, q2 = self.cri.get_q1_q2(state, action) # twin critics
obj_critic = ((self.criterion(q1, q_label) + self.criterion(q2, q_label)) * is_weights).mean()
td_error = (q_label - torch.min(q1, q1).detach()).abs()
buffer.td_error_update(td_error)
return obj_critic, state
class AgentSAC(AgentBase):
def __init__(self, net_dim, state_dim, action_dim, learning_rate=1e-4):
super().__init__()
self.target_entropy = np.log(action_dim)
self.alpha_log = torch.tensor(
(-np.log(action_dim) * np.e, ),
dtype=torch.float32,
requires_grad=True,
device=self.device) # trainable parameter
self.act = ActorSAC(net_dim, state_dim, action_dim).to(self.device)
self.act_target = deepcopy(self.act)
self.cri = CriticTwin(
int(net_dim * 1.25),
state_dim,
action_dim,
).to(self.device)
self.cri_target = deepcopy(self.cri)
self.criterion = torch.nn.SmoothL1Loss()
self.optimizer = torch.optim.Adam([{
'params': self.act.parameters(),
'lr': learning_rate * 0.75
}, {
'params': self.cri.parameters(),
'lr': learning_rate * 1.25
}, {
'params': (self.alpha_log, ),
'lr': learning_rate
}])
def select_actions(self, states): # states = (state, ...)
states = torch.as_tensor(states,
dtype=torch.float32,
device=self.device)
actions = self.act.get_action(states)
return actions.detach().cpu().numpy()
def update_policy(self, buffer, max_step, batch_size, repeat_times):
buffer.update__now_len__before_sample()
alpha = self.alpha_log.exp().detach()
obj_actor = obj_critic = None
for _ in range(int(max_step * repeat_times)):
with torch.no_grad():
reward, mask, action, state, next_s = buffer.random_sample(
batch_size)
next_a, next_log_prob = self.act_target.get__action__log_prob(
next_s)
next_q = torch.min(*self.cri_target.get__q1_q2(next_s, next_a))
q_label = reward + mask * (next_q + next_log_prob * alpha)
q1, q2 = self.cri.get__q1_q2(state, action)
obj_critic = self.criterion(q1, q_label) + self.criterion(
q2, q_label)
action_pg, log_prob = self.act.get__action__log_prob(
state) # policy gradient
obj_alpha = (self.alpha_log *
(log_prob - self.target_entropy).detach()).mean()
alpha = self.alpha_log.exp().detach()
with torch.no_grad():
self.alpha_log[:] = self.alpha_log.clamp(-16, 2)
obj_actor = -(
torch.min(*self.cri_target.get__q1_q2(state, action_pg)) +
log_prob * alpha).mean()
obj_united = obj_critic + obj_alpha + obj_actor
self.optimizer.zero_grad()
obj_united.backward()
self.optimizer.step()
soft_target_update(self.cri_target, self.cri)
soft_target_update(self.act_target, self.act)
# return obj_actor.item(), obj_critic.item()
return alpha.item(), obj_critic.item()