Replace strings with tuples in HeteroData keys. Refactor code for pre…

…training. Add pretraining dataset generation.

.gitignore

@@ -162,3 +162,4 @@ cython_debug/
 
 # MacOS specific
 .DS_Store
+AIAgent/report/

AIAgent/common/game.py

@@ -20,10 +20,10 @@ class State:
     VisitedAgainVertices: int
     VisitedNotCoveredVerticesInZone: int
     VisitedNotCoveredVerticesOutOfZone: int
-    StepWhenMovedLastTime: int
-    InstructionsVisitedInCurrentBlock: int
     History: list[StateHistoryElem]
     Children: list[int]
+    StepWhenMovedLastTime: int
+    InstructionsVisitedInCurrentBlock: int
 
     def __hash__(self) -> int:
         return self.Id.__hash__()

AIAgent/ml/common_model/paths.py

@@ -1,7 +1,6 @@
 import os
 import pathlib
 
-ROOT = "AIAgent"
 
 CSV_PATH = os.path.join("report", "epochs_tables")
 MODELS_PATH = os.path.join("report", "epochs_best")
@@ -12,7 +11,3 @@
 TRAINING_DATA_PATH = os.path.join("report", "run_tables")
 PRETRAINED_MODEL_PATH = os.path.join("ml", "models")
 RAW_FILES_PATH = os.path.join("report", "SerializedEpisodes")
-
-PATH_TO_MODELS_FOR_PARALLEL_ARCHITECTURE = os.path.join(
-    ROOT, "ml", "pretrained_models", "models_for_parallel_architecture"
-)

AIAgent/ml/data_loader_compact.py

@@ -21,7 +21,6 @@
 from functools import partial
 
 
-# NUM_NODE_FEATURES = 49
 NUM_NODE_FEATURES = 6
 EXPECTED_FILENAME = "expectedResults.txt"
 GAMESTATESUFFIX = "_gameState"
@@ -67,7 +66,7 @@ def __init__(self, raw_files_path: Path, processed_files_path: Path):
         self.processed_files_path = processed_files_path
 
     @staticmethod
-    def convert_input_to_tensor(input: GameState) -> Tuple[HeteroData, StateMap]:
+    def convert_input_to_tensor(input: GameState) -> Tuple[HeteroData, Dict[int, int]]:
         """
         Converts game env to tensors
         """
@@ -91,8 +90,8 @@ def convert_input_to_tensor(input: GameState) -> Tuple[HeteroData, StateMap]:
         edges_attr_s_v = []
         edges_attr_v_s = []
 
-        state_map: StateMap = {}  # Maps real state id to its index
-        vertex_map: VertexMap = {}  # Maps real vertex id to its index
+        state_map: Dict[int, int] = {}  # Maps real state id to its index
+        vertex_map: Dict[int, int] = {}  # Maps real vertex id to its index
         vertex_index = 0
         state_index = 0
 
@@ -110,6 +109,8 @@ def convert_input_to_tensor(input: GameState) -> Tuple[HeteroData, StateMap]:
                             int(v.CoveredByTest),
                             int(v.VisitedByState),
                             int(v.TouchedByState),
+                            int(v.ContainsCall),
+                            int(v.ContainsThrow),
                         ]
                     )
                 )
@@ -132,11 +133,12 @@ def convert_input_to_tensor(input: GameState) -> Tuple[HeteroData, StateMap]:
                     np.array(
                         [
                             s.Position,
-                            # s.PredictedUsefulness,
                             s.PathConditionSize,
                             s.VisitedAgainVertices,
                             s.VisitedNotCoveredVerticesInZone,
                             s.VisitedNotCoveredVerticesOutOfZone,
+                            s.StepWhenMovedLastTime,
+                            s.InstructionsVisitedInCurrentBlock,
                         ]
                     )
                 )
@@ -145,16 +147,24 @@ def convert_input_to_tensor(input: GameState) -> Tuple[HeteroData, StateMap]:
                     v_to = vertex_map[h.GraphVertexId]
                     edges_index_s_v_history.append(np.array([state_index, v_to]))
                     edges_index_v_s_history.append(np.array([v_to, state_index]))
-                    edges_attr_s_v.append(np.array([h.NumOfVisits]))
-                    edges_attr_v_s.append(np.array([h.NumOfVisits]))
+                    edges_attr_s_v.append(
+                        np.array([h.NumOfVisits, h.StepWhenVisitedLastTime])
+                    )
+                    edges_attr_v_s.append(
+                        np.array([h.NumOfVisits, h.StepWhenVisitedLastTime])
+                    )
                 state_index = state_index + 1
             else:
                 state_doubles += 1
 
         # state and its childen edges: state -> state
         for s in game_states:
             for ch in s.Children:
-                edges_index_s_s.append(np.array([state_map[s.Id], state_map[ch]]))
+                try:
+                    edges_index_s_s.append(np.array([state_map[s.Id], state_map[ch]]))
+                except KeyError:
+                    print("[ERROR]: KeyError")
+                    return None, None
 
         # state position edges: vertex -> state and back
         for v in graphVertices:
@@ -164,108 +174,67 @@ def convert_input_to_tensor(input: GameState) -> Tuple[HeteroData, StateMap]:
 
         data["game_vertex"].x = torch.tensor(np.array(nodes_vertex), dtype=torch.float)
         data["state_vertex"].x = torch.tensor(np.array(nodes_state), dtype=torch.float)
-        data["game_vertex_to_game_vertex"].edge_index = (
+
+        def tensor_not_empty(tensor):
+            return tensor.numel() != 0
+
+        # dumb fix
+        def null_if_empty(tensor):
+            return (
+                tensor
+                if tensor_not_empty(tensor)
+                else torch.empty((2, 0), dtype=torch.int64)
+            )
+
+        data["game_vertex", "to", "game_vertex"].edge_index = null_if_empty(
             torch.tensor(np.array(edges_index_v_v), dtype=torch.long).t().contiguous()
         )
-        data["game_vertex_to_game_vertex"].edge_attr = torch.tensor(
+
+        data["game_vertex", "to", "game_vertex"].edge_attr = torch.tensor(
             np.array(edges_attr_v_v), dtype=torch.long
         )
-        data["game_vertex_to_game_vertex"].edge_type = torch.tensor(
+        data["game_vertex", "to", "game_vertex"].edge_type = torch.tensor(
             np.array(edges_types_v_v), dtype=torch.long
         )
-        data["state_vertex_in_game_vertex"].edge_index = (
+        data["state_vertex", "in", "game_vertex"].edge_index = (
             torch.tensor(np.array(edges_index_s_v_in), dtype=torch.long)
             .t()
             .contiguous()
         )
-        data["game_vertex_in_state_vertex"].edge_index = (
+        data["game_vertex", "in", "state_vertex"].edge_index = (
             torch.tensor(np.array(edges_index_v_s_in), dtype=torch.long)
             .t()
             .contiguous()
         )
 
-        def tensor_not_empty(tensor):
-            return tensor.numel() != 0
-
-        # dumb fix
-        def null_if_empty(tensor):
-            return (
-                tensor
-                if tensor_not_empty(tensor)
-                else torch.empty((2, 0), dtype=torch.int64)
-            )
-
-        data["state_vertex_history_game_vertex"].edge_index = null_if_empty(
+        data["state_vertex", "history", "game_vertex"].edge_index = null_if_empty(
             torch.tensor(np.array(edges_index_s_v_history), dtype=torch.long)
             .t()
             .contiguous()
         )
-        data["game_vertex_history_state_vertex"].edge_index = null_if_empty(
+        data["game_vertex", "history", "state_vertex"].edge_index = null_if_empty(
             torch.tensor(np.array(edges_index_v_s_history), dtype=torch.long)
             .t()
             .contiguous()
         )
-        data["state_vertex_history_game_vertex"].edge_attr = torch.tensor(
+        data["state_vertex", "history", "game_vertex"].edge_attr = torch.tensor(
             np.array(edges_attr_s_v), dtype=torch.long
         )
-        data["game_vertex_history_state_vertex"].edge_attr = torch.tensor(
+        data["game_vertex", "history", "state_vertex"].edge_attr = torch.tensor(
             np.array(edges_attr_v_s), dtype=torch.long
         )
         # if (edges_index_s_s): #TODO: empty?
-        data["state_vertex_parent_of_state_vertex"].edge_index = null_if_empty(
+        data["state_vertex", "parent_of", "state_vertex"].edge_index = null_if_empty(
             torch.tensor(np.array(edges_index_s_s), dtype=torch.long).t().contiguous()
         )
-        # print(data['state', 'parent_of', 'state'].edge_index)
-        # data['game_vertex', 'to', 'game_vertex'].edge_attr = torch.tensor(np.array(edges_attr_v_v), dtype=torch.long)
-        # data['state_vertex', 'to', 'game_vertex'].edge_attr = torch.tensor(np.array(edges_attr_s_v), dtype=torch.long)
-        # data.state_map = state_map
-        # print("Doubles", state_doubles, len(state_map))
         return data, state_map
 
-    def process_directory(self, data_dir):
-        example_dirs = next(walk(data_dir), (None, [], None))[1]
-        example_dirs.sort()
-        print(example_dirs)
-        for fldr in example_dirs:
-            fldr_path = os.path.join(data_dir, fldr)
-            graphs_to_convert = []
-            for f in os.listdir(fldr_path):
-                if GAMESTATESUFFIX in f:
-                    graphs_to_convert.append(f)
-            graphs_to_convert.sort(key=lambda x: int(x.split("_")[0]))
-            self.graph_types_and_data[fldr] = graphs_to_convert
-
-    def __process_files(self):
-        for k, v in self.graph_types_and_data.items():
-            for file in v:
-                with open(os.path.join(self.data_dir, k, file)) as f:
-                    print(os.path.join(self.data_dir, k, file))
-                    data = json.load(f)
-                    graph, state_map = self.convert_input_to_tensor(
-                        GameState.from_dict(data)
-                    )
-                    if graph is not None:
-                        # add_expected values
-                        expected = self.get_expected_value(
-                            os.path.join(
-                                self.data_dir, k, file.split("_")[0] + STATESUFFIX
-                            ),
-                            state_map,
-                        )
-                        # print(len(graph['state_vertex'].x), len(state_map), len(expected))
-                        graph.y = expected
-                        self.dataset.append(graph)
-            PIK = "./dataset_t/" + k + ".dat"
-            with open(PIK, "wb") as f:
-                pickle.dump(self.dataset, f)
-            self.dataset = []
-
     def process_step(self, map_path: Path, step_id: str) -> Step:
         def get_states_properties(
             file_path: str, state_map: StateMap
         ) -> Step.StatesProperties:
             """Get tensor for states"""
-            expected = {}
+            expected = dict()
             with open(file_path) as f:
                 data = json.load(f)
                 state_set = set()
@@ -284,6 +253,7 @@ def get_states_properties(
                             d["ExpectedWeight"],
                         ]
                         expected[sid] = np.array(values)
+                f.close()
             ordered = []
             ordered_by_index = list(
                 zip(*sorted(state_map.items(), key=lambda x: x[1]))
@@ -295,7 +265,7 @@ def get_states_properties(
         def get_states_distribution(
             file_path: str, state_map: StateMap
         ) -> Step.StatesDistribution:
-            states_distribution = torch.zeros([len(state_map.keys())])
+            states_distribution = torch.zeros([len(state_map.keys()), 1])
             with open(file_path) as f:
                 state_id = int(f.read())
             states_distribution[state_map[state_id]] = 1
@@ -346,14 +316,25 @@ def get_step_ids(map_name: str):
                 )
 
     def save_dataset_for_training(self, dataset_state_path: Path, num_processes=1):
+        with open(dataset_state_path, "w") as csv_file:
+            writer = csv.writer(csv_file)
+            for map_stat in self.load_dataset(num_processes):
+                steps = []
+                for step in map_stat.Steps:
+                    step.Graph.y_true = step.StatesDistribution
+                    steps.append(step.Graph)
+                torch.save(
+                    steps,
+                    os.path.join(self.processed_files_path, map_stat.MapName + ".pt"),
+                )
+                writer.writerow([map_stat.MapName, map_stat.Result])
+
+    def save_dataset_for_pretraining(self, num_processes=1):
         for map_stat in self.load_dataset(num_processes):
             steps = []
             for step in map_stat.Steps:
-                step.Graph.y_true = step.StatesDistribution
+                step.Graph.y = step.StatesProperties
                 steps.append(step.Graph)
-            torch.save(
-                steps, os.path.join(self.processed_files_path, map_stat.MapName + ".pt")
-            )
-            with open(dataset_state_path, "a") as csv_file:
-                writer = csv.writer(csv_file)
-                writer.writerow([map_stat.MapName, map_stat.Result])
+            PIK = os.path.join(self.processed_files_path, map_stat.MapName + ".dat")
+            with open(PIK, "wb") as f:
+                pickle.dump(steps, f)

AIAgent/ml/het_gnn_test_train.py

@@ -5,6 +5,8 @@
 import torch
 import torch.nn.functional as F
 from torch_geometric.loader import DataLoader
+import tqdm
+from config import GeneralConfig
 
 BALANCE_DATASET = False
 
@@ -22,8 +24,6 @@ def __init__(self, model_class, hidden):
         self.hidden = hidden
 
     def train_and_save(self, dataset_dir, epochs, dir_to_save):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(device)
         dataset = []
         for file in os.listdir(dataset_dir):
             print(file)
@@ -46,17 +46,19 @@ def train_and_save(self, dataset_dir, epochs, dir_to_save):
         print(f"Number of test graphs: {len(test_dataset)}")
 
         train_loader = DataLoader(
-            train_dataset, batch_size=1, shuffle=False
+            train_dataset, batch_size=100, shuffle=False
         )  # TODO: add learning by batches!
         test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
 
-        model = self.model_class(hidden_channels=self.hidden, out_channels=8).to(device)
+        model = self.model_class(hidden_channels=self.hidden, out_channels=8).to(
+            GeneralConfig.DEVICE
+        )
         optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
 
         for epoch in range(1, epochs + 1):
-            self.train(model, train_loader, optimizer, device)
-            train_acc = self.tst(model, train_loader)
-            test_acc = self.tst(model, test_loader)
+            self.train(model, train_loader, optimizer, GeneralConfig.DEVICE)
+            train_acc = self.tst(model, train_loader, GeneralConfig.DEVICE)
+            test_acc = self.tst(model, test_loader, GeneralConfig.DEVICE)
             print(
                 f"Epoch: {epoch:03d}, Train Loss: {train_acc:.6f}, Test Loss: {test_acc:.6f}"
             )
@@ -70,39 +72,42 @@ def weighted_mse_loss(self, pred, target, weight=None):
 
     def train(self, model, train_loader, optimizer, device):
         model.train()
-        for data in train_loader:  # Iterate in batches over the training dataset.
+        for data in tqdm.tqdm(
+            train_loader, desc="Pretraining"
+        ):  # Iterate in batches over the training dataset.
             data = data.to(device)
             optimizer.zero_grad()  # Clear gradients.
             out = model(
                 data.x_dict["game_vertex"],
                 data.x_dict["state_vertex"],
-                data.edge_index_dict["game_vertex_to_game_vertex"],
-                data["game_vertex_to_game_vertex"].edge_type,
-                data["game_vertex_history_state_vertex"].edge_index,
-                data["game_vertex_history_state_vertex"].edge_attr,
-                data["game_vertex_in_state_vertex"].edge_index,
-                data["state_vertex_parent_of_state_vertex"].edge_index,
+                data.edge_index_dict["game_vertex", "to", "game_vertex"],
+                data["game_vertex", "to", "game_vertex"].edge_type,
+                data["game_vertex", "history", "state_vertex"].edge_index,
+                data["game_vertex", "history", "state_vertex"].edge_attr,
+                data["game_vertex", "in", "state_vertex"].edge_index,
+                data["state_vertex", "parent_of", "state_vertex"].edge_index,
             )
             target = data.y
             loss = F.mse_loss(out, target)
             loss.backward()  # Derive gradients.
             optimizer.step()  # Update parameters based on gradients.
 
     @torch.no_grad()
-    def tst(self, model, loader):
+    def tst(self, model, loader, device):
         model.eval()
         total_loss = 0
         number_of_states_total = 0
-        for data in loader:
+        for data in tqdm.tqdm(loader, desc="Testing"):
+            data.to(device)
             out = model(
                 data.x_dict["game_vertex"],
                 data.x_dict["state_vertex"],
-                data["game_vertex_to_game_vertex"].edge_index,
-                data["game_vertex_to_game_vertex"].edge_type,
-                data["game_vertex_history_state_vertex"].edge_index,
-                data["game_vertex_history_state_vertex"].edge_attr,
-                data["game_vertex_in_state_vertex"].edge_index,
-                data["state_vertex_parent_of_state_vertex"].edge_index,
+                data["game_vertex", "to", "game_vertex"].edge_index,
+                data["game_vertex", "to", "game_vertex"].edge_type,
+                data["game_vertex", "history", "state_vertex"].edge_index,
+                data["game_vertex", "history", "state_vertex"].edge_attr,
+                data["game_vertex", "in", "state_vertex"].edge_index,
+                data["state_vertex", "parent_of", "state_vertex"].edge_index,
             )
             target = data.y
             for i, x in enumerate(out):