minor changes

main.py

@@ -22,13 +22,13 @@ parser.add_argument('--tau', type=float, default=0.005, metavar='G',
                     help='target smoothing coefficient(τ) (default: 0.005)')
 parser.add_argument('--lr', type=float, default=0.0003, metavar='G',
                     help='learning rate (default: 0.0003)')
-parser.add_argument('--alpha', type=float, default=0.1, metavar='G',
-                    help='Temperature parameter α determines the relative importance of the entropy term against the reward (default: 0.1)')
+parser.add_argument('--alpha', type=float, default=0.2, metavar='G',
+                    help='Temperature parameter α determines the relative importance of the entropy term against the reward (default: 0.2)')
 parser.add_argument('--automatic_entropy_tuning', type=bool, default=False, metavar='G',
                     help='Temperature parameter α automaically adjusted.')
 parser.add_argument('--seed', type=int, default=456, metavar='N',
                     help='random seed (default: 456)')
-parser.add_argument('--batch_size', type=int, default=256, metavar='N',
+parser.add_argument('--batch_size', type=int, default=100, metavar='N',
                     help='batch size (default: 256)')
 parser.add_argument('--num_steps', type=int, default=1000001, metavar='N',
                     help='maximum number of steps (default: 1000000)')
@@ -48,9 +48,9 @@ args = parser.parse_args()
 
 # Environment
 env = NormalizedActions(gym.make(args.env_name))
-env.seed(args.seed)
 torch.manual_seed(args.seed)
 np.random.seed(args.seed)
+env.seed(args.seed)
 
 # Agent
 agent = SAC(env.observation_space.shape[0], env.action_space, args)
@@ -66,18 +66,15 @@ test_rewards = []
 total_numsteps = 0
 updates = 0
 
-for i_episode in itertools.count():
+for i_episode in itertools.count(1):
     state = env.reset()
-
     episode_reward = 0
+
     while True:
         if args.start_steps > total_numsteps:
             action = env.action_space.sample()
         else:
             action = agent.select_action(state)  # Sample action from policy
-        next_state, reward, done, _ = env.step(action)  # Step
-        mask = not done  # 1 for not done and 0 for done
-        memory.push(state, action, reward, next_state, mask)  # Append transition to memory
         if len(memory) > args.batch_size:
             for i in range(args.updates_per_step): # Number of updates per step in environment
                 # Sample a batch from memory
@@ -95,6 +92,11 @@ for i_episode in itertools.count():
                 writer.add_scalar('entropy_temprature/alpha', alpha, updates)
                 updates += 1
 
+        next_state, reward, done, _ = env.step(action)  # Step
+        mask = float(not done)  # 1 for not done and 0 for done
+
+        memory.push(state, action, reward, next_state, mask) # Append transition to memory
+
         state = next_state
         total_numsteps += 1
         episode_reward += reward

model.py

@@ -11,8 +11,7 @@ epsilon = 1e-6
 
 # Initialize Policy weights
 def weights_init_(m):
-    classname = m.__class__.__name__
-    if classname.find('Linear') != -1:
+    if isinstance(m, nn.Linear):
         torch.nn.init.xavier_uniform_(m.weight, gain=1)
         torch.nn.init.constant_(m.bias, 0)
 
@@ -51,13 +50,13 @@ class QNetwork(nn.Module):
         self.apply(weights_init_)
 
     def forward(self, state, action):
-        x1 = torch.cat([state, action], 1)
-        x1 = F.relu(self.linear1(x1))
+        xu = torch.cat([state, action], 1)
+        
+        x1 = F.relu(self.linear1(xu))
         x1 = F.relu(self.linear2(x1))
         x1 = self.linear3(x1)
 
-        x2 = torch.cat([state, action], 1)
-        x2 = F.relu(self.linear4(x2))
+        x2 = F.relu(self.linear4(xu))
         x2 = F.relu(self.linear5(x2))
         x2 = self.linear6(x2)
 
@@ -65,9 +64,10 @@ class QNetwork(nn.Module):
 
 
 class GaussianPolicy(nn.Module):
-    def __init__(self, num_inputs, num_actions, hidden_dim):
+    def __init__(self, num_inputs, num_actions, hidden_dim, max_act):
         super(GaussianPolicy, self).__init__()
-
+        
+        self.max_action = max_act
         self.linear1 = nn.Linear(num_inputs, hidden_dim)
         self.linear2 = nn.Linear(hidden_dim, hidden_dim)
 
@@ -94,7 +94,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, x_t, mean, log_std
+        return action, log_prob, torch.tanh(mean)
 
 class DeterministicPolicy(nn.Module):
     def __init__(self, num_inputs, num_actions, hidden_dim):

sac.py

@@ -12,6 +12,7 @@ class SAC(object):
     def __init__(self, num_inputs, action_space, args):
 
         self.num_inputs = num_inputs
+        self.max_action = float(action_space.high[0])
         self.action_space = action_space.shape[0]
         self.gamma = args.gamma
         self.tau = args.tau
@@ -32,11 +33,9 @@ class SAC(object):
                 self.target_entropy = -torch.prod(torch.Tensor(action_space.shape).to(self.device)).item()
                 self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
                 self.alpha_optim = Adam([self.log_alpha], lr=args.lr)
-            else:
-                pass
 
 
-            self.policy = GaussianPolicy(self.num_inputs, self.action_space, args.hidden_size).to(self.device)
+            self.policy = GaussianPolicy(self.num_inputs, self.action_space, args.hidden_size, self.max_action).to(self.device)
             self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)
 
             self.value = ValueNetwork(self.num_inputs, args.hidden_size).to(self.device)
@@ -56,15 +55,10 @@ class SAC(object):
         state = torch.FloatTensor(state).to(self.device).unsqueeze(0)
         if eval == False:
             self.policy.train()
-            action, _, _, _, _ = self.policy.sample(state)
+            action, _, _ = self.policy.sample(state)
         else:
             self.policy.eval()
-            _, _, _, action, _ = self.policy.sample(state)
-            if self.policy_type == "Gaussian":
-                action = torch.tanh(action)
-            else:
-                pass
-        #action = torch.tanh(action)
+            _, _, action = self.policy.sample(state)
         action = action.detach().cpu().numpy()
         return action[0]
 
@@ -75,95 +69,61 @@ class SAC(object):
         next_state_batch = torch.FloatTensor(next_state_batch).to(self.device)
         action_batch = torch.FloatTensor(action_batch).to(self.device)
         reward_batch = torch.FloatTensor(reward_batch).to(self.device).unsqueeze(1)
-        mask_batch = torch.FloatTensor(np.float32(mask_batch)).to(self.device).unsqueeze(1)
+        mask_batch = torch.FloatTensor(mask_batch).to(self.device).unsqueeze(1)
 
-        """
-        Use two Q-functions to mitigate positive bias in the policy improvement step that is known
-        to degrade performance of value based methods. Two Q-functions also significantly speed
-        up training, especially on harder task.
-        """
-        expected_q1_value, expected_q2_value = self.critic(state_batch, action_batch)
-        new_action, log_prob, _, mean, log_std = self.policy.sample(state_batch)
+        qf1, qf2 = self.critic(state_batch, action_batch) # Two Q-functions to mitigate positive bias in the policy improvement step
+        pi, log_pi, _ = self.policy.sample(state_batch)
 
         if self.policy_type == "Gaussian":
             if self.automatic_entropy_tuning:
-                """
-                Alpha Loss
-                """
-                alpha_loss = -(self.log_alpha * (log_prob + self.target_entropy).detach()).mean()
+                alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
                 self.alpha_optim.zero_grad()
                 alpha_loss.backward()
                 self.alpha_optim.step()
                 self.alpha = self.log_alpha.exp()
-                alpha_logs = self.alpha.clone() # For TensorboardX logs
+                alpha_logs = torch.tensor(self.alpha) # For TensorboardX logs
             else:
                 alpha_loss = torch.tensor(0.).to(self.device)
-                alpha_logs = self.alpha # For TensorboardX logs
+                alpha_logs = torch.tensor(self.alpha) # For TensorboardX logs
 
-
-            """
-            Including a separate function approximator for the soft value can stabilize training.
-            """
-            expected_value = self.value(state_batch)
-            target_value = self.value_target(next_state_batch)
-            next_q_value = reward_batch + mask_batch * self.gamma * (target_value).detach()
+            vf = self.value(state_batch) # separate function approximator for the soft value can stabilize training.
+            with torch.no_grad():
+                vf_next_target = self.value_target(next_state_batch)
+                next_q_value = reward_batch + mask_batch * self.gamma * (vf_next_target)
         else:
-            """
-            There is no need in principle to include a separate function approximator for the state value.
-            We use a target critic network for deterministic policy and eradicate the value value network completely.
-            """
             alpha_loss = torch.tensor(0.).to(self.device)
             alpha_logs = self.alpha  # For TensorboardX logs
-            next_state_action, _, _, _, _, = self.policy.sample(next_state_batch)
-            target_critic_1, target_critic_2 = self.critic_target(next_state_batch, next_state_action)
-            target_critic = torch.min(target_critic_1, target_critic_2)
-            next_q_value = reward_batch + mask_batch * self.gamma * (target_critic).detach()
+            with torch.no_grad():
+                next_state_action, _, _, _, _, = self.policy.sample(next_state_batch)
+                # Use a target critic network for deterministic policy and eradicate the value value network completely.
+                qf1_next_target, qf2_next_target = self.critic_target(next_state_batch, next_state_action)
+                min_qf_next_target = torch.min(qf1_next_target, qf2_next_target)
+                next_q_value = reward_batch + mask_batch * self.gamma * (min_qf_next_target)
         
         
-        """
-        Soft Q-function parameters can be trained to minimize the soft Bellman residual
-        JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
-        ∇JQ = ∇Q(st,at)(Q(st,at) - r(st,at) - γV(target)(st+1))
-        """
-        q1_value_loss = F.mse_loss(expected_q1_value, next_q_value)
-        q2_value_loss = F.mse_loss(expected_q2_value, next_q_value)
-        q1_new, q2_new = self.critic(state_batch, new_action)
-        expected_new_q_value = torch.min(q1_new, q2_new)
+        qf1_loss = F.mse_loss(qf1, next_q_value) # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
+        qf2_loss = F.mse_loss(qf2, next_q_value) # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
+        qf1_pi, qf2_pi = self.critic(state_batch, pi)
+        min_qf_pi = torch.min(qf1_pi, qf2_pi)
 
         if self.policy_type == "Gaussian":
-            """
-            Including a separate function approximator for the soft value can stabilize training and is convenient to 
-            train simultaneously with the other networks
-            Update the V towards the min of two Q-functions in order to reduce overestimation bias from function approximation error.
-            JV = 𝔼st~D[0.5(V(st) - (𝔼at~π[Qmin(st,at) - α * log π(at|st)]))^2]
-            ∇JV = ∇V(st)(V(st) - Q(st,at) + (α * logπ(at|st)))
-            """
-            next_value = expected_new_q_value - (self.alpha * log_prob)
-            value_loss = F.mse_loss(expected_value, next_value.detach())
-        else:
-            pass
+            vf_target = min_qf_pi - (self.alpha * log_pi)
+            value_loss = F.mse_loss(vf, vf_target.detach()) # JV = 𝔼st~D[0.5(V(st) - (𝔼at~π[Qmin(st,at) - α * log π(at|st)]))^2]
 
-        """
-        Reparameterization trick is used to get a low variance estimator
-        f(εt;st) = action sampled from the policy
-        εt is an input noise vector, sampled from some fixed distribution
-        Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]
-        ∇Jπ = ∇log π + ([∇at (α * logπ(at|st)) − ∇at Q(st,at)])∇f(εt;st)
-        """
-        policy_loss = ((self.alpha * log_prob) - expected_new_q_value).mean()
+        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
-        mean_loss = 0.001 * mean.pow(2).mean()
-        std_loss = 0.001 * log_std.pow(2).mean()
+        # mean_loss = 0.001 * mean.pow(2).mean()
+        # std_loss = 0.001 * log_std.pow(2).mean()
 
-        policy_loss += mean_loss + std_loss
+        # policy_loss += mean_loss + std_loss
 
         self.critic_optim.zero_grad()
-        q1_value_loss.backward()
+        qf1_loss.backward()
         self.critic_optim.step()
 
         self.critic_optim.zero_grad()
-        q2_value_loss.backward()
+        qf2_loss.backward()
         self.critic_optim.step()
 
         if self.policy_type == "Gaussian":
@@ -187,7 +147,7 @@ class SAC(object):
 
         elif updates % self.target_update_interval == 0 and self.policy_type == "Gaussian":
             soft_update(self.value_target, self.value, self.tau)
-        return value_loss.item(), q1_value_loss.item(), q2_value_loss.item(), policy_loss.item(), alpha_loss.item(), alpha_logs
+        return value_loss.item(), qf1_loss.item(), qf2_loss.item(), policy_loss.item(), alpha_loss.item(), alpha_logs.item()
 
     # Save model parameters
     def save_model(self, env_name, suffix="", actor_path=None, critic_path=None, value_path=None):