Refactor logger injection in rollout runners to improve logging flexibility and enable custom logger usage.

(Issue  - null)

Author: EnliteAI Bot

maze/core/rollout/rollout_runner.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import logging
 import os
 import time
 import traceback
@@ -27,9 +26,6 @@
 import numpy as np
 from omegaconf import DictConfig, OmegaConf
 
-logger = logging.getLogger('SEQUENTIAL RUNNER')
-logger.setLevel(logging.INFO)
-
 
 class RolloutRunner(Runner, ABC):
     """General abstract class for rollout runners.
@@ -241,6 +237,7 @@ def run_interaction_loop(
         env_seeds: list[Any],
         agent_seeds: list[Any],
         deterministic: bool,
+        logger: Logger,
         render: bool = False,
         after_reset_callback: Callable = None,
     ) -> None:
@@ -252,6 +249,8 @@ def run_interaction_loop(
         :param n_episodes: Count of episodes to perform.
         :param env_seeds: The env seeds to be used for each episode.
         :param agent_seeds: The agent seeds to be used for each episode.
+        :param deterministic: Argmax policy.
+        :param logger: Logger instance for logging rollout events.
         :param render: Whether to render the environment after every step.
         :param after_reset_callback: If supplied, this will be executed after each episode to notify the observer.
         """

maze/core/rollout/sequential_rollout_runner.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import logging
+
 from maze.core.annotations import override
 from maze.core.log_events.log_events_writer_registry import LogEventsWriterRegistry
 from maze.core.log_events.log_events_writer_tsv import LogEventsWriterTSV
@@ -21,6 +23,9 @@
 
 from tqdm import tqdm
 
+logger = logging.getLogger('SEQUENTIAL RUNNER')
+logger.setLevel(logging.INFO)
+
 
 class SequentialRolloutRunner(RolloutRunner):
     """Runs rollout in the local process. Useful for short rollouts or debugging.
@@ -107,6 +112,7 @@ def run_with(self, env: ConfigType, wrappers: CollectionOfConfigType, agent: Con
             env_seeds=env_seeds,
             agent_seeds=agent_seeds,
             deterministic=self.deterministic,
+            logger=logger,
         )
         self.progress_bar.close()
         env.write_epoch_stats()