refactor(log): use dataclass for config

scripts/train.py

@@ -36,13 +36,9 @@ def main(dict_cfg: DictConfig):
     cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True)
     init_logger(
         run_name=run_name,
-        config=cfg_dict,
-        project_name=config.logging.wandb_project_name,
-        entity=config.logging.wandb_entity,
+        full_config=cfg_dict,
+        logging_cfg=config.logging,
         base_dir=os.path.dirname(run_dir),
-        use_wandb=config.logging.track,
-        upload_final_model=config.logging.upload_final_model,
-        upload_checkpoints=config.logging.upload_checkpoints,
     )
     logger = get_logger()
     logger.info(f"Hydra-initialized run: {run_name}")

src/experiment_logger/simple_logger.py

@@ -14,18 +14,16 @@ class SimpleLogger:
     def __init__(
         self,
         run_name: str = "simple_run",
-        config: Optional[Dict[str, Any]] = None,
-        project_name: str = "none",
-        entity: Optional[str] = None,
+        full_config: Optional[Dict[str, Any]] = None,
+        logging_cfg: Optional[Any] = None,
         base_dir: str = "runs",
-        use_wandb: bool = False,
         save_code: bool = False,
         log_level: int = logging.INFO,
         _set_as_global: bool = False,
     ):
         self.is_interactive = True
         self.run_name = run_name
-        self.config = config or {}
+        self.full_config = full_config or {}
         print(f"[INIT] SimpleLogger initialized for run: {run_name}")
 
     def set_level(self, level: int):

src/experiment_logger/unified_logger.py

@@ -18,6 +18,7 @@ import jax.numpy as jnp
 import numpy as np
 
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
+from experiment_logger.config_logger import LoggingConfig
 
 # Global storage for the active logger and the proxy singleton
 _active_logger: Optional[Any] = None
@@ -88,13 +89,9 @@ class UnifiedLogger:
     def __init__(
         self,
         run_name: str,
-        config: Dict[str, Any],
-        project_name: str = "PPO-Modularity",
-        entity: Optional[str] = None,
+        full_config: Dict[str, Any],
+        logging_cfg: LoggingConfig,
         base_dir: str = "runs",
-        use_wandb: bool = True,
-        upload_final_model: bool = False,
-        upload_checkpoints: bool = False,
         save_code: bool = True,
         log_level: int = logging.INFO,
     ):
@@ -102,18 +99,16 @@ class UnifiedLogger:
 
         Args:
             run_name: Unique name for this run
-            config: Configuration dictionary with hyperparameters
-            project_name: WandB project name
-            entity: WandB entity (team/user name)
+            full_config: Full configuration dictionary with hyperparameters to be saved
+            logging_cfg: Structured logging configuration dataclass
             base_dir: Base directory for local storage
-            use_wandb: Whether to use WandB logging
             save_code: Whether to save code to WandB
         """
         self.run_name = run_name
-        self.config = config
-        self.use_wandb = use_wandb
-        self.upload_final_model = upload_final_model
-        self.upload_checkpoints = upload_checkpoints
+        self.full_config = full_config
+        self.use_wandb = logging_cfg.track
+        self.upload_final_model = logging_cfg.upload_final_model
+        self.upload_checkpoints = logging_cfg.upload_checkpoints
         self.wandb_available = False
         self.wandb_run = None
         self.is_interactive = sys.stdout.isatty()
@@ -163,7 +158,7 @@ class UnifiedLogger:
 
         # Initialize WandB if requested
         if self.use_wandb:
-            self._init_wandb(project_name, entity, save_code)
+            self._init_wandb(logging_cfg.wandb_project_name, logging_cfg.wandb_entity, save_code)
 
         # Initialize metrics storage
         self.metrics_buffer: List[Dict[str, Any]] = []
@@ -211,7 +206,7 @@ class UnifiedLogger:
             project=project_name,
             entity=entity,
             name=self.run_name,
-            config=self.config,
+            config=self.full_config,
             save_code=save_code,
             resume="allow",
         )
@@ -221,7 +216,7 @@ class UnifiedLogger:
         """Save configuration to disk."""
         try:
             with open(self.config_file, "w") as f:
-                yaml.dump(self.config, f, default_flow_style=False, indent=2, sort_keys=False)
+                yaml.dump(self.full_config, f, default_flow_style=False, indent=2, sort_keys=False)
             self.info(f"Config saved to {self.config_file}")
         except Exception as e:
             self.error(f"Error saving config: {e}")