#917 - Integrated the PrimaiteSession into all tests.

- Ran a full pre-commit hook and thus encountered tons of fixes required

src/primaite/agents/rllib.py

@@ -1,28 +1,35 @@
+from __future__ import annotations
+
 import json
 from datetime import datetime
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Union
 
+import tensorflow as tf
 from ray.rllib.algorithms import Algorithm
 from ray.rllib.algorithms.a2c import A2CConfig
 from ray.rllib.algorithms.ppo import PPOConfig
 from ray.tune.logger import UnifiedLogger
 from ray.tune.registry import register_env
-import tensorflow as tf
+
 from primaite import getLogger
 from primaite.agents.agent import AgentSessionABC
-from primaite.common.enums import AgentFramework, AgentIdentifier, \
-    DeepLearningFramework
+from primaite.common.enums import (
+    AgentFramework,
+    AgentIdentifier,
+    DeepLearningFramework,
+)
 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"],
         lay_down_config_path=env_config["lay_down_config_path"],
         session_path=env_config["session_path"],
-        timestamp_str=env_config["timestamp_str"]
+        timestamp_str=env_config["timestamp_str"],
     )
 
 
@@ -37,16 +44,15 @@ def _custom_log_creator(session_path: Path):
 
 
 class RLlibAgent(AgentSessionABC):
+    """An AgentSession class that implements a Ray RLlib agent."""
 
-    def __init__(
-            self,
-            training_config_path,
-            lay_down_config_path
-    ):
+    def __init__(self, training_config_path, lay_down_config_path):
         super().__init__(training_config_path, lay_down_config_path)
         if not self._training_config.agent_framework == AgentFramework.RLLIB:
-            msg = (f"Expected RLLIB agent_framework, "
-                   f"got {self._training_config.agent_framework}")
+            msg = (
+                f"Expected RLLIB agent_framework, "
+                f"got {self._training_config.agent_framework}"
+            )
             _LOGGER.error(msg)
             raise ValueError(msg)
         if self._training_config.agent_identifier == AgentIdentifier.PPO:
@@ -54,8 +60,10 @@ class RLlibAgent(AgentSessionABC):
         elif self._training_config.agent_identifier == AgentIdentifier.A2C:
             self._agent_config_class = A2CConfig
         else:
-            msg = ("Expected PPO or A2C agent_identifier, "
-                   f"got {self._training_config.agent_identifier.value}")
+            msg = (
+                "Expected PPO or A2C agent_identifier, "
+                f"got {self._training_config.agent_identifier.value}"
+            )
             _LOGGER.error(msg)
             raise ValueError(msg)
         self._agent_config: PPOConfig
@@ -86,8 +94,12 @@ class RLlibAgent(AgentSessionABC):
             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"]
+        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}")
@@ -106,43 +118,48 @@ class RLlibAgent(AgentSessionABC):
                 training_config_path=self._training_config_path,
                 lay_down_config_path=self._lay_down_config_path,
                 session_path=self.session_path,
-                timestamp_str=self.timestamp_str
-            )
+                timestamp_str=self.timestamp_str,
+            ),
         )
 
         self._agent_config.training(
             train_batch_size=self._training_config.num_steps
         )
-        self._agent_config.framework(
-            framework="tf"
-        )
+        self._agent_config.framework(framework="tf")
 
         self._agent_config.rollouts(
             num_rollout_workers=1,
             num_envs_per_worker=1,
-            horizon=self._training_config.num_steps
+            horizon=self._training_config.num_steps,
         )
         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
         episode_count = self._current_result["episodes_total"]
         if checkpoint_n > 0 and episode_count > 0:
-            if (
-                    (episode_count % checkpoint_n == 0)
-                    or (episode_count == self._training_config.num_episodes)
+            if (episode_count % checkpoint_n == 0) or (
+                episode_count == self._training_config.num_episodes
             ):
-                self._agent.save(self.checkpoints_path)
+                self._agent.save(str(self.checkpoints_path))
 
     def learn(
-            self,
-            time_steps: Optional[int] = None,
-            episodes: Optional[int] = None,
-            **kwargs
+        self,
+        time_steps: Optional[int] = None,
+        episodes: Optional[int] = None,
+        **kwargs,
     ):
+        """
+        Evaluate the agent.
+
+        :param time_steps: The number of steps per episode. Optional. If not
+            passed, the value from the training config will be used.
+        :param episodes: The number of episodes. Optional. If not
+            passed, the value from the training config will be used.
+        :param kwargs: Any agent-specific key-word args to be passed.
+        """
         # Temporarily override train_batch_size and horizon
         if time_steps:
             self._agent_config.train_batch_size = time_steps
@@ -150,37 +167,53 @@ class RLlibAgent(AgentSessionABC):
 
         if not episodes:
             episodes = self._training_config.num_episodes
-        _LOGGER.info(f"Beginning learning for {episodes} episodes @"
-                     f" {time_steps} time steps...")
+        _LOGGER.info(
+            f"Beginning learning for {episodes} episodes @"
+            f" {time_steps} time steps..."
+        )
         for i in range(episodes):
             self._current_result = self._agent.train()
             self._save_checkpoint()
-        if self._training_config.deep_learning_framework != DeepLearningFramework.TORCH:
+        if (
+            self._training_config.deep_learning_framework
+            != DeepLearningFramework.TORCH
+        ):
             policy = self._agent.get_policy()
             tf.compat.v1.summary.FileWriter(
-                self.session_path / "ray_results",
-                policy.get_session().graph
+                self.session_path / "ray_results", policy.get_session().graph
             )
         super().learn()
         self._agent.stop()
 
     def evaluate(
-            self,
-            time_steps: Optional[int] = None,
-            episodes: Optional[int] = None,
-            **kwargs
+        self,
+        time_steps: Optional[int] = None,
+        episodes: Optional[int] = None,
+        **kwargs,
     ):
+        """
+        Evaluate the agent.
+
+        :param time_steps: The number of steps per episode. Optional. If not
+            passed, the value from the training config will be used.
+        :param episodes: The number of episodes. Optional. If not
+            passed, the value from the training config will be used.
+        :param kwargs: Any agent-specific key-word args to be passed.
+        """
         raise NotImplementedError
 
     def _get_latest_checkpoint(self):
         raise NotImplementedError
 
     @classmethod
-    def load(cls):
+    def load(cls, path: Union[str, Path]) -> RLlibAgent:
+        """Load an agent from file."""
         raise NotImplementedError
 
     def save(self):
+        """Save the agent."""
         raise NotImplementedError
 
     def export(self):
+        """Export the agent to transportable file format."""
         raise NotImplementedError