diff --git a/model.py b/model.py index 3acbe8b..315a882 100644 --- a/model.py +++ b/model.py @@ -91,7 +91,7 @@ class GaussianPolicy(nn.Module): # Enforcing Action Bound log_prob -= torch.log(1 - action.pow(2) + epsilon) log_prob = log_prob.sum(1, keepdim=True) - return action, log_prob, torch.tanh(mean) + return action, log_prob, mean, log_std class DeterministicPolicy(nn.Module): def __init__(self, num_inputs, num_actions, hidden_dim): diff --git a/sac.py b/sac.py index dc86478..ab9416f 100644 --- a/sac.py +++ b/sac.py @@ -34,9 +34,10 @@ class SAC(object): def select_action(self, state, eval=False): state = torch.FloatTensor(state).to(self.device).unsqueeze(0) if eval == False: - action, _, _ = self.policy.sample(state) + action, _, _, _ = self.policy.sample(state) else: - _, _, action = self.policy.sample(state) + _, _, action, _ = self.policy.sample(state) + action = torch.tanh(action) action = action.detach().cpu().numpy() return action[0] @@ -68,12 +69,15 @@ class SAC(object): qf2_loss.backward() self.critic_optim.step() - pi, log_pi, _ = self.policy.sample(state_batch) + pi, log_pi, mean, log_std = self.policy.sample(state_batch) qf1_pi, qf2_pi = self.critic(state_batch, pi) min_qf_pi = torch.min(qf1_pi, qf2_pi) policy_loss = ((self.alpha * log_pi) - min_qf_pi).mean() # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))] + # Regularization Loss + reg_loss = 0.001 * (mean.pow(2).mean() + log_std.pow(2).mean()) + policy_loss += reg_loss vf = self.value(state_batch) with torch.no_grad():