#917 - Almost there. All output files being writen for SB3/RLLIB PPO & A2C. Just need to bring in the hardcoded agents then update the testa and docs.

src/primaite/agents/rllib.py

@@ -1,18 +1,23 @@
-import glob
-import time
-from enum import Enum
+import json
+from datetime import datetime
 from pathlib import Path
-from typing import Union, Optional
+from pathlib import Path
+from typing import Optional
 
 from ray.rllib.algorithms import Algorithm
 from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.algorithms.a2c import A2CConfig
+from ray.tune.logger import UnifiedLogger
 from ray.tune.registry import register_env
 
+from primaite import getLogger
 from primaite.agents.agent import AgentSessionABC
-from primaite.config import training_config
+from primaite.common.enums import AgentFramework, RedAgentIdentifier
 from primaite.environment.primaite_env import Primaite
 
 
+_LOGGER = getLogger(__name__)
+
 def _env_creator(env_config):
     return Primaite(
         training_config_path=env_config["training_config_path"],
@@ -23,7 +28,17 @@ def _env_creator(env_config):
     )
 
 
-class RLlibPPO(AgentSessionABC):
+def _custom_log_creator(session_path: Path):
+    logdir = session_path / "ray_results"
+    logdir.mkdir(parents=True, exist_ok=True)
+
+    def logger_creator(config):
+        return UnifiedLogger(config, logdir, loggers=None)
+
+    return logger_creator
+
+
+class RLlibAgent(AgentSessionABC):
 
     def __init__(
             self,
@@ -31,17 +46,63 @@ class RLlibPPO(AgentSessionABC):
             lay_down_config_path
     ):
         super().__init__(training_config_path, lay_down_config_path)
-        self._ppo_config: PPOConfig
+        if not self._training_config.agent_framework == AgentFramework.RLLIB:
+            msg = (f"Expected RLLIB agent_framework, "
+                   f"got {self._training_config.agent_framework}")
+            _LOGGER.error(msg)
+            raise ValueError(msg)
+        if self._training_config.red_agent_identifier == RedAgentIdentifier.PPO:
+            self._agent_config_class = PPOConfig
+        elif self._training_config.red_agent_identifier == RedAgentIdentifier.A2C:
+            self._agent_config_class = A2CConfig
+        else:
+            msg = ("Expected PPO or A2C red_agent_identifier, "
+                   f"got {self._training_config.red_agent_identifier.value}")
+            _LOGGER.error(msg)
+            raise ValueError(msg)
+        self._agent_config: PPOConfig
+
         self._current_result: dict
         self._setup()
-        self._agent.save()
+        _LOGGER.debug(
+            f"Created {self.__class__.__name__} using: "
+            f"agent_framework={self._training_config.agent_framework}, "
+            f"red_agent_identifier="
+            f"{self._training_config.red_agent_identifier}, "
+            f"deep_learning_framework="
+            f"{self._training_config.deep_learning_framework}"
+        )
+
+    def _update_session_metadata_file(self):
+        """
+        Update the ``session_metadata.json`` file.
+
+        Updates the `session_metadata.json`` in the ``session_path`` directory
+        with the following key/value pairs:
+
+        - end_datetime: The date & time the session ended in iso format.
+        - total_episodes: The total number of training episodes completed.
+        - total_time_steps: The total number of training time steps completed.
+        """
+        with open(self.session_path / "session_metadata.json", "r") as file:
+            metadata_dict = json.load(file)
+
+        metadata_dict["end_datetime"] = datetime.now().isoformat()
+        metadata_dict["total_episodes"] = self._current_result["episodes_total"]
+        metadata_dict["total_time_steps"] = self._current_result["timesteps_total"]
+
+        filepath = self.session_path / "session_metadata.json"
+        _LOGGER.debug(f"Updating Session Metadata file: {filepath}")
+        with open(filepath, "w") as file:
+            json.dump(metadata_dict, file)
+            _LOGGER.debug("Finished updating session metadata file")
 
     def _setup(self):
         super()._setup()
         register_env("primaite", _env_creator)
-        self._ppo_config = PPOConfig()
+        self._agent_config = self._agent_config_class()
 
-        self._ppo_config.environment(
+        self._agent_config.environment(
             env="primaite",
             env_config=dict(
                 training_config_path=self._training_config_path,
@@ -52,19 +113,21 @@ class RLlibPPO(AgentSessionABC):
             )
         )
 
-        self._ppo_config.training(
+        self._agent_config.training(
             train_batch_size=self._training_config.num_steps
         )
-        self._ppo_config.framework(
-            framework=self._training_config.deep_learning_framework.value
+        self._agent_config.framework(
+            framework=self._training_config.deep_learning_framework
         )
 
-        self._ppo_config.rollouts(
+        self._agent_config.rollouts(
             num_rollout_workers=1,
             num_envs_per_worker=1,
             horizon=self._training_config.num_steps
         )
-        self._agent: Algorithm = self._ppo_config.build()
+        self._agent: Algorithm = self._agent_config.build(
+            logger_creator=_custom_log_creator(self.session_path)
+        )
 
     def _save_checkpoint(self):
         checkpoint_n = self._training_config.checkpoint_every_n_episodes
@@ -84,8 +147,8 @@ class RLlibPPO(AgentSessionABC):
     ):
         # Temporarily override train_batch_size and horizon
         if time_steps:
-            self._ppo_config.train_batch_size = time_steps
-            self._ppo_config.horizon = time_steps
+            self._agent_config.train_batch_size = time_steps
+            self._agent_config.horizon = time_steps
 
         if not episodes:
             episodes = self._training_config.num_episodes