feat: main and train CLI args

Author: jacazek

FlashAttentionTransformer_weights.pth

@@ -27,7 +27,7 @@ def generate_pong_states(num_steps=None):
 @inject.params(game_state=State)
 def _generate_pong_states(game_state: State = None):
     dt = 1  # Time step
-    ball_random_velocity = random_velocity_generator()
+    ball_random_velocity = random_velocity_generator(min=game_state.engineConfig.min_ball_velocity, max=game_state.engineConfig.max_ball_velocity)
 
     left_paddle = game_state.left_paddle
     right_paddle = game_state.right_paddle

RNNModel_weights.pth

@@ -6,8 +6,9 @@
 import mlflow
 
 from game.state import State
+from main_arguments import MainArguments
 from models import ModelConfiguration
-from runtime_configuration import mlflow_model_path, model_path, classification_threshold, temperature_variance, mlflow_server_url
+from runtime_configuration import mlflow_model_path, classification_threshold, temperature_variance, mlflow_server_url
 import numpy as np
 from model_loaders import load_mlflow_model, load_pytorch_model
 
@@ -26,7 +27,7 @@
     print(f"Failed to connect to MLflow server at {mlflow_server_url}. Error: {e}")
     print("Will load models from local mlruns directory")
 
-config = ModelConfiguration()
+# config = ModelConfiguration()
 
 def generate_fuzzy_states(num_steps=None):
     state_generator = _generate_fuzzy_states()
@@ -37,33 +38,36 @@ def generate_fuzzy_states(num_steps=None):
         for step in range(num_steps):
             yield next(state_generator)
 
-@inject.params(game_state=State)
-def _generate_fuzzy_states(game_state=State):
+@inject.params(game_state=State, main_arguments=MainArguments)
+def _generate_fuzzy_states(game_state=State, main_arguments=MainArguments):
     dt = 1  # Time step
-
     # Either load model from mlflow run
     # model = load_mlflow_model(mlflow_model_path)
 
     # Or load the model from pth file containing weights
-    model = load_pytorch_model(model_path)
+    model = load_pytorch_model(f"{main_arguments.model_type}_weights.pth")
+
 
     model.eval()
-    window_size = config.input_sequence_length
+    window_size = main_arguments.input_sequence_length
     window = deque(maxlen=window_size)
-    window.extend(np.zeros((config.input_sequence_length, config.input_size), dtype=float))
+    window.extend(np.zeros((main_arguments.input_sequence_length, main_arguments.input_size), dtype=float))
     while True:
         game_state.left_paddle.update(dt)
         game_state.right_paddle.update(dt)
         paddle_data = game_state.left_paddle.vectorize_state() + game_state.right_paddle.vectorize_state()
-        temperature = torch.from_numpy(
-            np.random.uniform(1.0 - temperature_variance, 1.0 + temperature_variance, config.discrete_output_size)).to(
-            device=config.device)
-        ball_data, discrete_data = model(torch.tensor(np.array([window])).to(device=config.device, dtype=torch.float), temperature)
+        # temperature = torch.from_numpy(
+        #     np.random.uniform(1.0 - temperature_variance, 1.0 + temperature_variance, main_arguments.discrete_output_size) * 100).to(
+        #     device=main_arguments.device)
+        temperature = 1 # larger temperature is more creativw
+        ball_data, discrete_data = model(torch.tensor(np.array([window])).to(device=main_arguments.device, dtype=torch.float), temperature)
+
         ball_data = ball_data.tolist()[0]
         discrete_probabilities = torch.sigmoid(discrete_data)
 
         classes = (discrete_probabilities > classification_threshold).int()
         classes = classes.tolist()[0]
+
         window.append(ball_data + paddle_data + classes)
         yield ball_data, paddle_data, classes[:4], classes[4:]
 

TransformerModel_weights.pth

@@ -1,6 +1,9 @@
 class EngineConfig:
     def __init__(self, ball_radius_percent=.01, paddle_width_percent=.01,
-                 paddle_height_percent=.2):
+                 paddle_height_percent=.2, min_ball_velocity=.005, max_ball_velocity=.025):
         self.ball_radius_percent = ball_radius_percent
         self.paddle_width_percent = paddle_width_percent
         self.paddle_height_percent = paddle_height_percent
+        self.max_ball_velocity = max_ball_velocity
+        self.min_ball_velocity = min_ball_velocity
+

exact_engine.py

@@ -11,6 +11,8 @@
 from fuzzy_engine import generate_fuzzy_states
 import pygame
 import inject
+from models import ModelConfiguration
+from main_arguments import MainArguments
 
 # Initialize Pygame
 pygame.init()
@@ -100,12 +102,11 @@ def render_state(state, count, engine_config: EngineConfig = None, field: Field
 
 # Main loop to render the state
 
-@inject.params(generator="generator")
-def main(generator):
+@inject.params(main_arguments=MainArguments)
+def main(main_arguments: MainArguments):
     global screen, screen_width, screen_height
     running = True
-    # for index, state in enumerate(generate_fuzzy_states()):
-    for index, state in enumerate(generator()):
+    for index, state in enumerate(main_arguments.generator()):
         if not running:
             break
         for event in pygame.event.get():
@@ -131,6 +132,9 @@ def main(generator):
     pygame.quit()
 
 def configure_main(binder: inject.Binder):
+    main_arguments = MainArguments.get_arguments()
+    binder.bind(MainArguments, main_arguments)
+    binder.bind(ModelConfiguration, main_arguments)
     # immediatly construct and bind an instance to the given key
     binder.bind(Field, Field(1.0, 1.0))
     binder.bind(EngineConfig, EngineConfig())
@@ -147,10 +151,6 @@ def configure_main(binder: inject.Binder):
     binder.bind_to_constructor(State, State)
 
 
-    # Choose the kind of generator desired
-    # binder.bind("generator", generate_pong_states)
-    binder.bind("generator", generate_fuzzy_states)
-
 if __name__ == "__main__":
     inject.configure(configure_main)
     main()

fuzzy_engine.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+import argparse
+import os
+import subprocess
+from models.rnn import RNNModel
+from models.transformer import TransformerModel
+from models.transformer_flashattn import FlashAttentionTransformer
+from models import ModelConfiguration
+from exact_engine import generate_pong_states
+from fuzzy_engine import generate_fuzzy_states
+
+generators = {
+    "exact": generate_pong_states,
+    "fuzzy": generate_fuzzy_states,
+}
+
+model_dictionary = {RNNModel.__name__: RNNModel,
+          TransformerModel.__name__: TransformerModel,
+          FlashAttentionTransformer.__name__: FlashAttentionTransformer}
+model_names = list(model_dictionary.keys())
+
+@dataclass
+class MainArguments(ModelConfiguration):
+    mlflow_server_url: str = "https://localhost:8080"
+
+    model_type: str = RNNModel.__name__
+
+    generator_type: str = list(generators.keys())[1]
+    generator = list(generators.values())[1]
+
+    # keep this parameter last
+    command: str = ""
+
+    @staticmethod
+    def get_arguments():
+        parser = argparse.ArgumentParser(description="Main configuration")
+
+        parser.add_argument("--mlflow_server_url", type=str, default="http://localhost:8080", help="mlflow server url")
+        parser.add_argument("--model_type", type=str, default=model_names[0],help="The model type to train", choices=model_names)
+        parser.add_argument("--generator_type", type=str, default=list(generators.keys())[1],help="The generator type to train", choices=list(generators.keys()))
+        parser.add_argument("--input_size", type=int, default=16, help="The input size of the model")
+        parser.add_argument("--hidden_size", type=int, default=128, help="The hidden size of the model")
+        parser.add_argument("--num_layers", type=int, default=2, help="The number of layers of the model")
+        parser.add_argument("--number_heads", type=int, default=16, help="The number of heads of the model (transformer model only)")
+        parser.add_argument("--input_sequence_length", type=int, default=10,help="The length of the input sequence")
+
+        args = parser.parse_args()
+        main_arguments = MainArguments()
+        for key, value in vars(args).items():
+            setattr(main_arguments, key, value)
+
+        main_arguments.generator = generators.get(args.__dict__["generator_type"])
+
+        main_arguments.command = str(subprocess.run(["ps", "-p", f"{os.getpid()}", "-o", "args", "--no-headers"], capture_output=True,
+               text=True).stdout)
+
+        return main_arguments

game/configuration.py

@@ -1,9 +1,7 @@
 import torch
 import mlflow.pytorch
-from runtime_configuration import Model
 from models import ModelConfiguration
-
-config = ModelConfiguration()
+import inject
 
 
 def save_mlflow_model(model, path):
@@ -13,12 +11,14 @@ def save_pytorch_model(model, path):
     torch.save(model.state_dict(), path)
 
 
-def load_mlflow_model(path):
+@inject.params(config=ModelConfiguration)
+def load_mlflow_model(path, config: ModelConfiguration):
     model = mlflow.pytorch.load_model(path, map_location=torch.device(config.device))
     return model
 
-def load_pytorch_model(path):
+@inject.params(config=ModelConfiguration)
+def load_pytorch_model(path, config: ModelConfiguration):
     """Load weights into a pytorch model from the specified path to .pth file"""
-    model = Model().to(device=config.device)
+    model = config.model().to(device=config.device)
     model.load_state_dict(torch.load(path, weights_only=True, map_location=torch.device(config.device)))
     return model

main.py

@@ -1,6 +1,6 @@
 # only export public facing stuff from the package
 from .model_configuration import ModelConfiguration
 from .pong_dataset import PongDataset # should probably move to separate package
-from .rnn import RNNModel
-from .transformer import TransformerModel
+# from .rnn import RNNModel
+# from .transformer import TransformerModel
 # from .transformer_flashattn import FlashAttentionTransformer

main_arguments.py

@@ -2,14 +2,15 @@
 
 import torch
 from torch import nn as nn
-from . import ModelConfiguration
+import inject
+from models import ModelConfiguration
 
-config = ModelConfiguration()
 
 
 class BasePongModel(nn.Module, ABC):
-    def __init__(self):
+    def __init__(self, config: ModelConfiguration):
         super(BasePongModel, self).__init__()
+        self.config = config
         # Linear layer to expand input from 10 to 64 dimensions
         self.fc_feature_expansion = nn.Linear(config.input_size, config.hidden_size)
 

model_loaders.py

@@ -1,17 +1,21 @@
 # general parameters
 import torch
+import dataclasses
 
 
+@dataclasses.dataclass
 class ModelConfiguration:
-    def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        # model parameters
-        self.input_size = 16
-        self.hidden_size = 128
-        self.output_size = 4
-        self.discrete_output_size = 6
-        self.num_layers = 2
-        self.number_heads = 8
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # model parameters
+    input_size = 16
+    hidden_size = 128
+    output_size = 4
+    discrete_output_size = 6
+    num_layers = 2
+    number_heads = 16
 
-        # training parameters
-        self.input_sequence_length=20
+    # training parameters
+    input_sequence_length=20
+
+    # def get_model_path(self):
+    #     return f"{self.model.__name__}_weights.pth"

models/__init__.py

@@ -1,15 +1,12 @@
 from torch import nn as nn
 
 from .base_pong_model import BasePongModel
-from . import ModelConfiguration
-
-config = ModelConfiguration()
 
 
 class RNNModel(BasePongModel):
-    def __init__(self):
-        super(RNNModel, self).__init__()
-        self.lstm = nn.LSTM(config.hidden_size, config.hidden_size, config.num_layers, batch_first=True, dropout=0.2)
+    def __init__(self, model_config):
+        super(RNNModel, self).__init__(model_config)
+        self.lstm = nn.LSTM(self.config.hidden_size, self.config.hidden_size, self.config.num_layers, batch_first=True, dropout=0.2)
 
     def _forward(self, x):
         out, _ = self.lstm(x)

models/base_pong_model.py

@@ -4,22 +4,19 @@
 from .base_pong_model import BasePongModel
 from . import ModelConfiguration
 
-config = ModelConfiguration()
-
-
 class TransformerModel(BasePongModel):
-    def __init__(self):
-        super(TransformerModel, self).__init__()
+    def __init__(self, model_config: ModelConfiguration):
+        super(TransformerModel, self).__init__(model_config)
         # Consider using decoder only with flash attention
-        self.positional_encoding = nn.Parameter(torch.zeros(1, 100, config.hidden_size))
+        self.positional_encoding = nn.Parameter(torch.zeros(1, 100, model_config.hidden_size))
 
         # self.transformer = nn.TransformerEncoder(
         self.transformer_list = nn.ModuleList([nn.TransformerEncoderLayer(
-            d_model=config.hidden_size,
-            nhead=config.number_heads,
-            dim_feedforward=config.hidden_size,
+            d_model=self.config.hidden_size,
+            nhead=self.config.number_heads,
+            dim_feedforward=self.config.hidden_size,
             batch_first=True,
-        ) for _ in range(config.num_layers)])
+        ) for _ in range(self.config.num_layers)])
         #     num_layers=num_layers,
         # )
         # self.transformer = nn.TransformerDecoderLayer(

models/model_configuration.py

@@ -5,9 +5,7 @@
 from flash_attn.modules.mha import MHA
 from flash_attn import flash_attn_func
 from .base_pong_model import BasePongModel
-from . import ModelConfiguration
 
-config = ModelConfiguration()
 
 
 class Transformer(nn.Module):
@@ -47,7 +45,7 @@ def _generate_positional_encoding(max_seq_len, embed_dim):
 class TransformerLayer(nn.Module):
     def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
         super(TransformerLayer, self).__init__()
-        self.mha = MHA(embed_dim, num_heads, causal=True, use_flash_attn=True, return_residual=False)
+        self.mha = MHA(embed_dim, num_heads, causal=True, use_flash_attn=False, return_residual=False)
         # self.mha = MultiHeadAttention(embed_dim, num_heads)
         self.ffn = FeedForwardNetwork(embed_dim, ff_dim, dropout)
         self.norm1 = nn.LayerNorm(embed_dim)
@@ -112,7 +110,7 @@ def _forward(self, x: torch.Tensor):
         x = self.transformer(x)
         return x.mean(dim=1)
 
-    def __init__(self):
-        super(FlashAttentionTransformer, self).__init__()
-        self.transformer = Transformer(config.hidden_size, config.number_heads, config.num_layers, config.hidden_size, config.input_sequence_length, 0.2)
+    def __init__(self, model_config):
+        super(FlashAttentionTransformer, self).__init__(model_config)
+        self.transformer = Transformer(self.config.hidden_size, self.config.number_heads, self.config.num_layers, self.config.hidden_size, self.config.input_sequence_length, 0.2)
 

models/rnn.py

@@ -1,18 +1,15 @@
-from models import TransformerModel, RNNModel
-# from models import FlashAttentionTransformer
+
+
 
 mlflow_server_url = "http://localhost:8080"
 
-# model to use during training and inference
-Model = RNNModel
 
 # path to mlflow model
 # use a run that corresponds with the desired model type (rnn/transformer/flashtransformer)
-mlflow_model_path = f"runs:/62a7d1ead3564c379cbacbff4ef7ac55/model_e99"
-
-# path to pytorch model weights
-model_path = f"{Model.__name__}_weights.pth"
+mlflow_model_path = f"runs:/6d4a6cb5a09c420ca834fe16795b16a3/model_e99"
 
 
 classification_threshold = 0.5
-temperature_variance = 0.0
+temperature_variance = 0.7
+
+# class RuntimeArguments: