fix and refactor

VSharp.ML.AIAgent/learning/play_game.py

@@ -2,33 +2,24 @@
 from statistics import StatisticsError
 from time import perf_counter
 from typing import TypeAlias
-import random
-import copy
 
 import tqdm
 from func_timeout import FunctionTimedOut, func_set_timeout
 
-from torch_geometric.data import HeteroData
-
-
 from common.classes import GameResult, Map2Result
 from common.constants import TQDM_FORMAT_DICT
 from common.game import GameMap
 from common.utils import get_states
 from config import FeatureConfig, GeneralConfig
 from connection.broker_conn.socket_manager import game_server_socket_manager
 from connection.game_server_conn.connector import Connector
-from connection.game_server_conn.utils import MapsType, get_maps
+from connection.game_server_conn.utils import MapsType
 from learning.timer.resources_manager import manage_map_inference_times_array
 from learning.timer.stats import compute_statistics
 from learning.timer.utils import get_map_inference_times
 from ml.data_loader_compact import ServerDataloaderHeteroVector
 from ml.fileop import save_model
 from ml.model_wrappers.protocols import Predictor
-from torch_geometric.data import Dataset
-
-# from ray.experimental.tqdm_ray import tqdm
-
 
 TimeDuration: TypeAlias = float
 

VSharp.ML.AIAgent/ml/common_model/dataset.py

@@ -47,32 +47,29 @@ def load(self):
     def remove_similar_steps(self, map_steps):
         filtered_map_steps = []
         for step in map_steps:
-            if len(filtered_map_steps) != 0:
-                if step["y_true"].size() == filtered_map_steps[-1]["y_true"].size():
-                    cos_d = 1 - torch.sum(
-                        (step["y_true"] / torch.linalg.vector_norm(step["y_true"]))
-                        * (
-                            filtered_map_steps[-1]["y_true"]
-                            / torch.linalg.vector_norm(filtered_map_steps[-1]["y_true"])
-                        )
+            if (
+                len(filtered_map_steps) != 0
+                and step["y_true"].size() == filtered_map_steps[-1]["y_true"].size()
+            ):
+                cos_d = 1 - torch.sum(
+                    (step["y_true"] / torch.linalg.vector_norm(step["y_true"]))
+                    * (
+                        filtered_map_steps[-1]["y_true"]
+                        / torch.linalg.vector_norm(filtered_map_steps[-1]["y_true"])
+                    )
+                )
+                if (
+                    cos_d < 1e-7
+                    and step["game_vertex"]["x"].size()[0]
+                    == filtered_map_steps[-1]["game_vertex"]["x"].size()[0]
+                ):
+                    step.use_for_train = np.random.choice(
+                        [True, False],
+                        p=[
+                            self.similar_steps_save_prob,
+                            1 - self.similar_steps_save_prob,
+                        ],
                     )
-                    if (
-                        cos_d < 1e-7
-                        and step["game_vertex"]["x"].size()[0]
-                        == filtered_map_steps[-1]["game_vertex"]["x"].size()[0]
-                    ):
-                        if np.random.choice(
-                            [True, False],
-                            p=[
-                                self.similar_steps_save_prob,
-                                1 - self.similar_steps_save_prob,
-                            ],
-                        ):
-                            step.use_for_train = True
-                        else:
-                            step.use_for_train = False
-                    else:
-                        step.use_for_train = True
                 else:
                     step.use_for_train = True
             else:
@@ -83,13 +80,11 @@ def remove_similar_steps(self, map_steps):
     def filter_map_steps(self, map_steps):
         filtered_map_steps = []
         for step in map_steps:
-            if step["y_true"].size()[0] != 1:
-                if not step["y_true"].isnan().any():
-                    max_ind = torch.argmax(step["y_true"])
-                    step["y_true"] = torch.zeros_like(step["y_true"])
-                    step["y_true"][max_ind] = 1.0
-                    filtered_map_steps.append(step)
-
+            if step["y_true"].size()[0] != 1 and not step["y_true"].isnan().any():
+                max_ind = torch.argmax(step["y_true"])
+                step["y_true"] = torch.zeros_like(step["y_true"])
+                step["y_true"][max_ind] = 1.0
+                filtered_map_steps.append(step)
         return filtered_map_steps
 
     def get_plain_data(self):

VSharp.ML.AIAgent/ml/common_model/models.py

@@ -46,52 +46,57 @@ def __init__(
 
         self.mlp = MLP(hidden_channels, [1])
 
-    def forward(self, x_dict, edge_index_dict, edge_attr_dict):
+    def forward(
+        self,
+        game_x,
+        state_x,
+        edge_index_v_v,
+        edge_type_v_v,
+        edge_index_history_v_s,
+        edge_attr_history_v_s,
+        edge_index_in_v_s,
+        edge_index_s_s,
+    ):
         game_x = self.gv_layers[0](
-            x_dict["game_vertex"],
-            edge_index_dict[("game_vertex", "to", "game_vertex")],
+            game_x,
+            edge_index_v_v,
         ).relu()
         for layer in self.gv_layers[1:]:
             game_x = layer(
                 game_x,
-                edge_index_dict[("game_vertex", "to", "game_vertex")],
+                edge_index_v_v,
             ).relu()
 
         state_x = self.sv_layers[0](
-            x_dict["state_vertex"],
-            edge_index_dict[("state_vertex", "parent_of", "state_vertex")],
+            state_x,
+            edge_index_s_s,
         ).relu()
         for layer in self.sv_layers[1:]:
             state_x = layer(
                 state_x,
-                edge_index_dict[("state_vertex", "parent_of", "state_vertex")],
+                edge_index_s_s,
             ).relu()
 
         history_x = self.history1(
             (game_x, state_x),
-            edge_index_dict[("game_vertex", "history", "state_vertex")],
-            edge_attr_dict,
+            edge_index_history_v_s,
+            edge_attr_history_v_s,
             size=(game_x.size(0), state_x.size(0)),
         ).relu()
 
-        in_x = self.in1(
-            (game_x, history_x), edge_index_dict[("game_vertex", "in", "state_vertex")]
-        ).relu()
+        in_x = self.in1((game_x, history_x), edge_index_in_v_s).relu()
 
         state_x = self.sv_layers2[0](
             in_x,
-            edge_index_dict[("state_vertex", "parent_of", "state_vertex")],
+            edge_index_s_s,
         ).relu()
         for layer in self.sv_layers2[1:]:
             state_x = layer(
                 state_x,
-                edge_index_dict[("state_vertex", "parent_of", "state_vertex")],
+                edge_index_s_s,
             ).relu()
         x = self.mlp(state_x)
-        z_dict = {}
-        z_dict["state_vertex"] = x
-        z_dict["game_vertex"] = x_dict["game_vertex"]
-        return z_dict
+        return x
 
 
 class ParallelBlocks(torch.nn.Module):
@@ -100,14 +105,30 @@ def __init__(self, models_list, mlp_list):
         self.models_list = models_list
         self.mlp = MLP(len(models_list), mlp_list)
 
-    def forward(self, x_dict, edge_index_dict, edge_attr=None):
+    def forward(
+        self,
+        game_x,
+        state_x,
+        edge_index_v_v,
+        edge_type_v_v,
+        edge_index_history_v_s,
+        edge_attr_history_v_s,
+        edge_index_in_v_s,
+        edge_index_s_s,
+    ):
         results_list = []
         for model in self.models_list:
             results_list.append(
-                model(x_dict, edge_index_dict, edge_attr)["state_vertex"]
+                model(
+                    game_x,
+                    state_x,
+                    edge_index_v_v,
+                    edge_type_v_v,
+                    edge_index_history_v_s,
+                    edge_attr_history_v_s,
+                    edge_index_in_v_s,
+                    edge_index_s_s,
+                )
             )
-        z_dict = {}
         results_tensor = torch.cat(results_list, dim=1)
-        z_dict["state_vertex"] = self.mlp(results_tensor)
-        z_dict["game_vertex"] = x_dict["game_vertex"]
-        return z_dict
+        return results_tensor

VSharp.ML.AIAgent/ml/common_model/utils.py

@@ -14,7 +14,9 @@
     common_models_path,
 )
 from ml.utils import load_model
-from ml.models.TAGSageTeacher.model_modified import StateModelEncoderLastLayer
+from ml.models.StateGNNEncoderConvEdgeAttr.model_modified import (
+    StateModelEncoderLastLayer,
+)
 
 
 def euclidean_dist(y_pred, y_true):
@@ -58,26 +60,25 @@ def csv2best_models():
                     models_stat[map_names[i]] = int_row[i]
                 models.append((row[0], models_stat))
 
-            for map_name in map_names:
-                best_model = max(models, key=(lambda m: m[1][map_name]))
-                best_model_name = best_model[0]
-                best_model_score = best_model[1]
-                path_to_model = os.path.join(
-                    models_path,
-                    "epoch_" + str(epoch_num),
-                    best_model_name + ".pth",
-                )
-                ref_model = load_model(
-                    Path(path_to_model), model=StateModelEncoderLastLayer(32, 8)
-                )
-
-                ref_model.to(GeneralConfig.DEVICE)
-                best_models[map_name] = [
-                    ref_model,
-                    best_model_score[map_name],
-                    best_model_name,
-                ]
-            return best_models
+        for map_name in map_names:
+            best_model = max(models, key=(lambda m: m[1][map_name]))
+            best_model_name, best_model_score = best_model[0], best_model[1]
+            path_to_model = os.path.join(
+                models_path,
+                "epoch_" + str(epoch_num),
+                best_model_name + ".pth",
+            )
+            ref_model = load_model(
+                Path(path_to_model), model=StateModelEncoderLastLayer(32, 8)
+            )
+
+            ref_model.to(GeneralConfig.DEVICE)
+            best_models[map_name] = (
+                ref_model,
+                best_model_score[map_name],
+                best_model_name,
+            )
+        return best_models
 
 
 def back_prop(best_model, model, data, optimizer, criterion):
@@ -127,9 +128,7 @@ def load_best_models_dict(path):
                 )
                 ref_model.load_state_dict(torch.load(path_to_model))
                 ref_model.to(GeneralConfig.DEVICE)
-                best_models[row[0]][0] = ref_model
-                best_models[row[0]][1] = ast.literal_eval(row[2])
-                best_models[row[0]][2] = row[1]
+                best_models[row[0]] = (ref_model, ast.literal_eval(row[2]), row[1])
     return best_models
 
 
@@ -142,7 +141,8 @@ def load_dataset_state_dict(path):
     return dataset_state_dict
 
 
-def get_model(path_to_weights: Path, model: torch.nn.Module):
+def get_model(path_to_weights: Path, model: torch.nn.Module, random_seed: int):
+    np.random.seed(random_seed)
     weights = torch.load(path_to_weights)
     weights["lin_last.weight"] = torch.tensor(np.random.random([1, 8]))
     weights["lin_last.bias"] = torch.tensor(np.random.random([1]))

VSharp.ML.AIAgent/ml/models/TAGSageTeacher/model.py → VSharp.ML.AIAgent/ml/models/StateGNNEncoderConvEdgeAttr/model.py

@@ -7,8 +7,8 @@ class StateModelEncoder(torch.nn.Module):
     def __init__(self, hidden_channels, out_channels):
         super().__init__()
         self.conv1 = TAGConv(5, hidden_channels, 2)
-        self.conv2 = TAGConv(6, hidden_channels, 3)  # TAGConv
-        self.conv3 = GraphConv((-1, -1), hidden_channels)  # SAGEConv
+        self.conv2 = TAGConv(6, hidden_channels, 3)
+        self.conv3 = GraphConv((-1, -1), hidden_channels)
         self.conv32 = GraphConv((-1, -1), hidden_channels)
         self.conv4 = SAGEConv((-1, -1), hidden_channels)
         self.conv42 = SAGEConv((-1, -1), hidden_channels)

VSharp.ML.AIAgent/ml/models/TAGSageTeacher/model_modified.py → VSharp.ML.AIAgent/ml/models/StateGNNEncoderConvEdgeAttr/model_modified.py

@@ -1,7 +1,5 @@
 import torch
-from torch.nn import Linear
 from torch_geometric.nn import Linear
-
 from torch.nn.functional import softmax
 from .model import StateModelEncoder
 

VSharp.ML.AIAgent/ml/models/TAGSageSimple/model.py

@@ -7,8 +7,8 @@ class StateModelEncoder(torch.nn.Module):
     def __init__(self, hidden_channels, out_channels):
         super().__init__()
         self.conv1 = TAGConv(5, hidden_channels, 2)
-        self.conv2 = TAGConv(hidden_channels, hidden_channels, 3)  # TAGConv
-        self.conv3 = GraphConv((-1, -1), hidden_channels)  # SAGEConv
+        self.conv2 = TAGConv(hidden_channels, hidden_channels, 3)
+        self.conv3 = GraphConv((-1, -1), hidden_channels)
         self.conv4 = SAGEConv((-1, -1), hidden_channels)
         self.lin = Linear(hidden_channels, out_channels)
 

VSharp.ML.AIAgent/ml/models/TAGSageSimple/model_modified.py

@@ -1,8 +1,4 @@
-import torch
-from torch.nn import Linear
 from torch_geometric.nn import Linear
-
-from learning.timer.wrapper import timeit
 from torch.nn.functional import softmax
 from .model import StateModelEncoder