更新算法模版

projects/codes/common/models.py

@@ -5,7 +5,7 @@ Author: John
 Email: johnjim0816@gmail.com
 Date: 2021-03-12 21:14:12
 LastEditor: John
-LastEditTime: 2022-08-29 14:24:44
+LastEditTime: 2022-10-31 23:53:06
 Discription: 
 Environment: 
 '''
@@ -35,20 +35,65 @@ class ActorSoftmax(nn.Module):
     def __init__(self, input_dim, output_dim, hidden_dim=256):
         super(ActorSoftmax, self).__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
-        self.fc2 = nn.Linear(hidden_dim, output_dim)
-    def forward(self,state):
-        dist = F.relu(self.fc1(state))
-        dist = F.softmax(self.fc2(dist),dim=1)
-        return dist
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, output_dim)
+    def forward(self,x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        probs = F.softmax(self.fc3(x),dim=1)
+        return probs
+
+class ActorSoftmaxTanh(nn.Module):
+    def __init__(self, input_dim, output_dim, hidden_dim=256):
+        super(ActorSoftmaxTanh, self).__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, output_dim)
+    def forward(self,x):
+        x = F.tanh(self.fc1(x))
+        x = F.tanh(self.fc2(x))
+        probs = F.softmax(self.fc3(x),dim=1)
+        return probs
+class ActorNormal(nn.Module):
+    def __init__(self, n_states,n_actions, hidden_dim=256):
+        super(ActorNormal, self).__init__()
+        self.fc1 = nn.Linear(n_states, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, n_actions)
+        self.fc4 = nn.Linear(hidden_dim, n_actions)
+    def forward(self,x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        mu = torch.tanh(self.fc3(x))
+        sigma = F.softplus(self.fc4(x)) + 0.001 # avoid 0
+        return mu,sigma
+# class ActorSoftmax(nn.Module):
+#     def __init__(self,input_dim, output_dim,
+#             hidden_dim=256):
+#         super(ActorSoftmax, self).__init__()
+#         self.actor = nn.Sequential(
+#                 nn.Linear(input_dim, hidden_dim),
+#                 nn.ReLU(),
+#                 nn.Linear(hidden_dim, hidden_dim),
+#                 nn.ReLU(),
+#                 nn.Linear(hidden_dim, output_dim),
+#                 nn.Softmax(dim=-1)
+#         )
+#     def forward(self, state):
+#         probs = self.actor(state)
+#         dist = Categorical(probs)
+#         return dist
 class Critic(nn.Module):
     def __init__(self,input_dim,output_dim,hidden_dim=256):
         super(Critic,self).__init__()
         assert output_dim == 1 # critic must output a single value
         self.fc1 = nn.Linear(input_dim, hidden_dim)
-        self.fc2 = nn.Linear(hidden_dim, output_dim)
-    def forward(self,state):
-        value = F.relu(self.fc1(state))
-        value = self.fc2(value)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, output_dim)
+    def forward(self,x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        value = self.fc3(x)
         return value
 
 class ActorCriticSoftmax(nn.Module):
@@ -72,18 +117,18 @@ class ActorCriticSoftmax(nn.Module):
         return value, policy_dist
 
 class ActorCritic(nn.Module):
-    def __init__(self, n_states, n_actions, hidden_dim=256):
+    def __init__(self, input_dim, output_dim, hidden_dim=256):
         super(ActorCritic, self).__init__()
         self.critic = nn.Sequential(
-            nn.Linear(n_states, hidden_dim),
+            nn.Linear(input_dim, hidden_dim),
             nn.ReLU(),
             nn.Linear(hidden_dim, 1)
         )
         
         self.actor = nn.Sequential(
-            nn.Linear(n_states, hidden_dim),
+            nn.Linear(input_dim, hidden_dim),
             nn.ReLU(),
-            nn.Linear(hidden_dim, n_actions),
+            nn.Linear(hidden_dim, output_dim),
             nn.Softmax(dim=1),
         )