#917 - Synced with dev and added better logging

src/primaite/agents/agent.py

@@ -3,11 +3,11 @@ import time
 from abc import ABC, abstractmethod
 from datetime import datetime
 from pathlib import Path
-from typing import Optional, Final, Dict, Union, List
+from typing import Optional, Final, Dict, Union
 from uuid import uuid4
 
+import primaite
 from primaite import getLogger, SESSIONS_DIR
-from primaite.common.enums import OutputVerboseLevel
 from primaite.config import lay_down_config
 from primaite.config import training_config
 from primaite.config.training_config import TrainingConfig
@@ -141,14 +141,13 @@ class AgentSessionABC(ABC):
 
     @abstractmethod
     def _setup(self):
-        if self.output_verbose_level >= OutputVerboseLevel.INFO:
-            _LOGGER.info(
-                "Welcome to the Primary-level AI Training Environment "
-                "(PrimAITE)"
-            )
-            _LOGGER.debug(
-                f"The output directory for this agent is: {self.session_path}"
-            )
+        _LOGGER.info(
+            "Welcome to the Primary-level AI Training Environment "
+            f"(PrimAITE) (version: {primaite.__version__})"
+        )
+        _LOGGER.info(
+            f"The output directory for this session is: {self.session_path}"
+        )
         self._write_session_metadata_file()
         self._can_learn = True
         self._can_evaluate = False
@@ -165,6 +164,7 @@ class AgentSessionABC(ABC):
             **kwargs
     ):
         if self._can_learn:
+            _LOGGER.info("Finished learning")
             _LOGGER.debug("Writing transactions")
             self._update_session_metadata_file()
             self._can_evaluate = True
@@ -176,7 +176,7 @@ class AgentSessionABC(ABC):
             episodes: Optional[int] = None,
             **kwargs
     ):
-        pass
+        _LOGGER.info("Finished evaluation")
 
     @abstractmethod
     def _get_latest_checkpoint(self):
@@ -260,6 +260,7 @@ class HardCodedAgentSessionABC(AgentSessionABC):
                 # Introduce a delay between steps
                 time.sleep(self._training_config.time_delay / 1000)
         self._env.close()
+        super().evaluate()
 
     @classmethod
     def load(cls):

src/primaite/agents/rllib.py

@@ -152,7 +152,8 @@ 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...")
         for i in range(episodes):
             self._current_result = self._agent.train()
             self._save_checkpoint()

src/primaite/agents/sb3.py

@@ -69,8 +69,9 @@ class SB3Agent(AgentSessionABC):
                     (episode_count % checkpoint_n == 0)
                     or (episode_count == self._training_config.num_episodes)
             ):
-                self._agent.save(
-                    self.checkpoints_path / f"sb3ppo_{episode_count}.zip")
+                checkpoint_path = self.checkpoints_path / f"sb3ppo_{episode_count}.zip"
+                self._agent.save(checkpoint_path)
+                _LOGGER.debug(f"Saved agent checkpoint: {checkpoint_path}")
 
     def _get_latest_checkpoint(self):
         pass
@@ -86,6 +87,8 @@ class SB3Agent(AgentSessionABC):
 
         if not episodes:
             episodes = self._training_config.num_episodes
+        _LOGGER.info(f"Beginning learning for {episodes} episodes @"
+                     f" {time_steps} time steps...")
         for i in range(episodes):
             self._agent.learn(total_timesteps=time_steps)
             self._save_checkpoint()
@@ -106,6 +109,9 @@ class SB3Agent(AgentSessionABC):
         if not episodes:
             episodes = self._training_config.num_episodes
 
+        _LOGGER.info(f"Beginning evaluation for {episodes} episodes @"
+                     f" {time_steps} time steps...")
+
         for episode in range(episodes):
             obs = self._env.reset()
 

src/primaite/agents/utils.py

@@ -6,8 +6,8 @@ from primaite.common.enums import (
     NodeHardwareAction,
     NodeSoftwareAction,
     SoftwareState,
+    NodePOLType
 )
-from primaite.common.enums import NodePOLType
 
 
 def transform_action_node_readable(action):