#917 - Got RLlib fully training in PrimAITE. Started integrating the the other agents into the Session class

src/primaite/config/_package_data/training/training_config_main.yaml

@@ -1,26 +1,52 @@
 # Main Config File
 
-# Generic config values
-# Choose one of these (dependent on Agent being trained)
-# "STABLE_BASELINES3_PPO"
-# "STABLE_BASELINES3_A2C"
-# "GENERIC"
-agent_identifier: STABLE_BASELINES3_A2C
+# Sets which agent algorithm framework will be used:
+# Options are:
+# "SB3" (Stable Baselines3)
+# "RLLIB" (Ray RLlib)
+# "NONE" (Custom Agent)
+agent_framework: RLLIB
+
+# Sets which deep learning framework will be used. Default is TF (Tensorflow).
+# Options are:
+# "TF" (Tensorflow)
+# TF2 (Tensorflow 2.X)
+# TORCH (PyTorch)
+deep_learning_framework: TORCH
+
+# Sets which Red Agent algo/class will be used:
+# Options are:
+# "A2C" (Advantage Actor Critic)
+# "PPO" (Proximal Policy Optimization)
+# "HARDCODED" (Custom Agent)
+# "RANDOM" (Random Action)
+red_agent_identifier: PPO
+
 # Sets How the Action Space is defined:
 # "NODE"
 # "ACL"
 # "ANY" node and acl actions
 action_type: NODE
+
 # Number of episodes to run per session
 num_episodes: 10
+
 # Number of time_steps per episode
 num_steps: 256
+
+# Sets how often the agent will save a checkpoint (every n time episodes).
+# Set to 0 if no checkpoints are required. Default is 10
+checkpoint_every_n_episodes: 5
+
 # Time delay between steps (for generic agents)
 time_delay: 10
+
 # Type of session to be run (TRAINING or EVALUATION)
 session_type: TRAINING
+
 # Determine whether to load an agent from file
 load_agent: False
+
 # File path and file name of agent if you're loading one in
 agent_load_file: C:\[Path]\[agent_saved_filename.zip]
 

src/primaite/config/training_config.py

@@ -8,6 +8,7 @@ from typing import Any, Dict, Final, Union, Optional
 import yaml
 
 from primaite import USERS_CONFIG_DIR, getLogger
+from primaite.common.enums import DeepLearningFramework
 from primaite.common.enums import ActionType, RedAgentIdentifier, \
     AgentFramework, SessionType
 
@@ -20,10 +21,13 @@ _EXAMPLE_TRAINING: Final[Path] = USERS_CONFIG_DIR / "example_config" / "training
 class TrainingConfig:
     """The Training Config class."""
     agent_framework: AgentFramework = AgentFramework.SB3
-    "The agent framework."
+    "The AgentFramework"
+
+    deep_learning_framework: DeepLearningFramework = DeepLearningFramework.TF
+    "The DeepLearningFramework."
 
     red_agent_identifier: RedAgentIdentifier = RedAgentIdentifier.PPO
-    "The red agent/algo class."
+    "The RedAgentIdentifier.."
 
     action_type: ActionType = ActionType.ANY
     "The ActionType to use."
@@ -33,6 +37,10 @@ class TrainingConfig:
 
     num_steps: int = 256
     "The number of steps in an episode."
+
+    checkpoint_every_n_episodes: int = 5
+    "The agent will save a checkpoint every n episodes."
+
     observation_space: dict = field(
         default_factory=lambda: {"components": [{"name": "NODE_LINK_TABLE"}]}
     )
@@ -148,6 +156,7 @@ class TrainingConfig:
     ) -> TrainingConfig:
         field_enum_map = {
             "agent_framework": AgentFramework,
+            "deep_learning_framework": DeepLearningFramework,
              "red_agent_identifier": RedAgentIdentifier,
              "action_type": ActionType,
              "session_type": SessionType
@@ -155,7 +164,7 @@ class TrainingConfig:
 
         for field, enum_class in field_enum_map.items():
             if field in config_dict:
-                config_dict[field] = enum_class[field]
+                config_dict[field] = enum_class[config_dict[field]]
 
         return TrainingConfig(**config_dict)
 
@@ -219,7 +228,7 @@ def load(file_path: Union[str, Path],
                 )
                 _LOGGER.error(msg)
         try:
-            return TrainingConfig.from_dict(**config)
+            return TrainingConfig.from_dict(config)
         except TypeError as e:
             msg = (
                 f"Error when creating an instance of {TrainingConfig} "