update

codes/A2C/model.py

@@ -1,36 +1,36 @@
 #!/usr/bin/env python
 # coding=utf-8
 '''
-Author: John
+Author: JiangJi
 Email: johnjim0816@gmail.com
-Date: 2020-11-03 20:45:25
-LastEditor: John
-LastEditTime: 2021-03-20 17:41:33
+Date: 2021-05-03 21:38:54
+LastEditor: JiangJi
+LastEditTime: 2021-05-03 21:40:06
 Discription: 
 Environment: 
 '''
 import torch.nn as nn
+import torch.nn.functional as F
 from torch.distributions import Categorical
-
 class ActorCritic(nn.Module):
-    def __init__(self, state_dim, action_dim, hidden_dim=256):
+    def __init__(self, num_inputs, num_outputs, hidden_size, std=0.0):
         super(ActorCritic, self).__init__()
+        
         self.critic = nn.Sequential(
-            nn.Linear(state_dim, hidden_dim),
+            nn.Linear(num_inputs, hidden_size),
             nn.ReLU(),
-            nn.Linear(hidden_dim, 1)
+            nn.Linear(hidden_size, 1)
         )
         
         self.actor = nn.Sequential(
-            nn.Linear(state_dim, hidden_dim),
+            nn.Linear(num_inputs, hidden_size),
             nn.ReLU(),
-            nn.Linear(hidden_dim, action_dim),
+            nn.Linear(hidden_size, num_outputs),
             nn.Softmax(dim=1),
         )
         
     def forward(self, x):
         value = self.critic(x)
-        print(x)
         probs = self.actor(x)
         dist  = Categorical(probs)
         return dist, value