fix: model and train

Author: wangxt111

core

@@ -10,7 +10,7 @@ def __init__(self, input_channel_num, num_actions, extra_size):
         super().__init__()
         self.channels = input_channel_num
         self.embeddings = nn.ModuleList(
-            [nn.Embedding(9, 16) for _ in range(input_channel_num)])
+            [nn.Embedding(10, 16) for _ in range(input_channel_num)])
         self.conv1 = nn.Conv2d(64, 64, kernel_size=8, stride=4)
         self.conv2 = nn.Conv2d(64, 64, kernel_size=4, stride=2)
         self.bn = nn.BatchNorm2d(64)
@@ -50,7 +50,7 @@ def __init__(self, input_channel_num, num_actions, extra_size):
         super().__init__()
         self.channels = input_channel_num
         self.embeddings = nn.ModuleList(
-            [nn.Embedding(9, 16) for _ in range(input_channel_num)])
+            [nn.Embedding(10, 16) for _ in range(input_channel_num)])
 
         self.conv1 = nn.Conv2d(64, 64, kernel_size=8, stride=4)
         self.conv2 = nn.Conv2d(64, 64, kernel_size=4, stride=2)
@@ -88,7 +88,7 @@ def forward(self, x, y):
 
 # test the shape of the output
 if __name__ == "__main__":
-    rand_input = torch.rand(1, 4, 38, 38)
+    rand_input = torch.rand(1, 4, 41, 41)
     extra_input = torch.rand(1, 10)
     pacman_net = PacmanNet(4, 5, 10)
     res = pacman_net(rand_input, extra_input)

model.py

@@ -25,7 +25,7 @@
 # initialize networks
 policy_net_pacman = PacmanNet(4, 5, 40)
 target_net_pacman = PacmanNet(4, 5, 40)
-policy_net_pacman.load_state_dict(torch.load("pacman.pth"))
+# policy_net_pacman.load_state_dict(torch.load("pacman.pth"))
 target_net_pacman.load_state_dict(policy_net_pacman.state_dict())
 target_net_pacman.eval()
 
@@ -34,7 +34,7 @@
 
 policy_net_ghost = GhostNet(4, 5, 40)
 target_net_ghost = GhostNet(4, 5, 40)
-policy_net_ghost.load_state_dict(torch.load("ghost.pth"))
+# policy_net_ghost.load_state_dict(torch.load("ghost.pth"))
 target_net_ghost.load_state_dict(policy_net_ghost.state_dict())
 target_net_ghost.eval()
 
@@ -87,24 +87,24 @@ def state_dict_to_tensor(state_dict):
         board = np.array(board)
     size = board.shape[0]
     # print(board)
-    # pad board to 38x38
-    padding_num = 38 - size
+    # pad board to 41x41
+    padding_num = 41 - size
     board = np.pad(board, pad_width=(0, padding_num),
                    mode="constant", constant_values=0)
     # pacman position matrix
-    pacman_pos = np.zeros((38, 38))
+    pacman_pos = np.zeros((41, 41))
     if "pacman_pos" in state_dict:
         pacman_pos[state_dict["pacman_pos"][0] + padding_num][
             state_dict["pacman_pos"][1] + padding_num
         ] = 1
 
     # ghost position matrix
-    ghost_pos = np.zeros((38, 38))
+    ghost_pos = np.zeros((41, 41))
     if "ghost_pos" in state_dict:
         for ghost in state_dict["ghost_pos"]:
             ghost_pos[ghost[0] + padding_num][ghost[1] + padding_num] = 1
 
-    portal_pos = np.zeros((38, 38))
+    portal_pos = np.zeros((41, 41))
     if "portal" in state_dict:
         portal = state_dict["portal"]
         if portal[0] != -1 and portal[1] != -1: