refactor(log): improved logging workflow

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -34,6 +34,13 @@ class BrittleStarJaxEnvWrapper:
 
         self._action_rng = None
 
+        from experiment_logger import get_logger
+
+        self.logger = get_logger()
+        self.logger.info(
+            f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}"
+        )
+
     @property
     def backend(self):
         return self._backend
@@ -51,6 +58,7 @@ class BrittleStarJaxEnvWrapper:
         return self._env.observation_space
 
     def reset(self, seed: int = 0):
+        self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
         self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
         env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
         return self._vectorized_reset(rng=env_rngs)

src/brittle_star_project/environment/factory.py

@@ -98,9 +98,15 @@ class BrittleStarEnvFactory:
             case _:
                 raise ValueError(f"Unsupported task: {env_config.task}")
 
-        return env_class.from_morphology_and_arena(
+        env = env_class.from_morphology_and_arena(
             morphology=morphology,
             arena=arena,
             configuration=env_configuration,
             backend=backend.value,
         )
+
+        from experiment_logger import get_logger
+
+        get_logger().info(f"Created {env_config.task.value} env on backend {backend.value}")
+
+        return env

src/experiment_logger/__init__.py

@@ -5,8 +5,15 @@ This package provides a unified interface for logging to multiple backends
 """
 
 from experiment_logger.config_utils import load_yaml_config, merge_config_with_cli
-from experiment_logger.unified_logger import UnifiedLogger
+from experiment_logger.unified_logger import UnifiedLogger, get_logger
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
 
-__all__ = ["UnifiedLogger", "init_wandb", "finish_wandb", "load_yaml_config", "merge_config_with_cli"]
+__all__ = [
+    "UnifiedLogger",
+    "get_logger",
+    "init_wandb",
+    "finish_wandb",
+    "load_yaml_config",
+    "merge_config_with_cli",
+]
 __version__ = "0.1.0"

src/experiment_logger/config_utils.py

@@ -1,15 +1,16 @@
 """Configuration utilities for loading YAML configs and merging with CLI args."""
 
-import logging
 import os
 import sys
 from typing import Dict, Any, Type, TypeVar
 import yaml
 from dataclasses import fields, is_dataclass
 
-log = logging.getLogger(__name__)
+from experiment_logger.unified_logger import get_logger
 
-T = TypeVar('T')
+log = get_logger()
+
+T = TypeVar("T")
 
 
 def load_yaml_config(config_path: str) -> Dict[str, Any]:
@@ -17,7 +18,7 @@ def load_yaml_config(config_path: str) -> Dict[str, Any]:
     if not os.path.exists(config_path):
         raise FileNotFoundError(f"Config file not found: {config_path}")
 
-    with open(config_path, 'r') as f:
+    with open(config_path, "r") as f:
         config = yaml.safe_load(f)
 
     if config is None:
@@ -31,7 +32,7 @@ def save_yaml_config(config: Dict[str, Any], config_path: str):
     """Save configuration to YAML file."""
     os.makedirs(os.path.dirname(config_path), exist_ok=True)
 
-    with open(config_path, 'w') as f:
+    with open(config_path, "w") as f:
         yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False)
 
     log.info(f"Saved configuration to: {config_path}")
@@ -55,13 +56,13 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
                 # Handle None values and optional types
                 if value is None:
                     filtered_config[key] = None
-                elif hasattr(field.type, '__origin__') and field.type.__origin__ is type(None):
+                elif hasattr(field.type, "__origin__") and field.type.__origin__ is type(None):
                     # Optional type (Union[X, None])
                     filtered_config[key] = value
                 else:
                     # Try to convert to the expected type
-                    if field.type == bool and isinstance(value, str):
+                    if field.type is bool and isinstance(value, str):
-                        filtered_config[key] = value.lower() in ('true', '1', 'yes', 'on')
+                        filtered_config[key] = value.lower() in ("true", "1", "yes", "on")
                     else:
                         filtered_config[key] = field.type(value) if value is not None else None
             except (ValueError, TypeError) as e:

src/experiment_logger/unified_logger.py

@@ -2,15 +2,16 @@
 
 This logger ensures all experimental data is preserved by writing to:
 1. Weights & Biases (when available)
-2. Local disk (JSON files, model checkpoints)
+2. Local disk (JSON files, model checkpoints, run.log)
 3. stdout (for real-time monitoring)
 """
 
 import json
 import logging
+import subprocess
 import time
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import flax
 import jax.numpy as jnp
@@ -18,7 +19,38 @@ import numpy as np
 
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
 
-logger = logging.getLogger(__name__)
+# Global singleton storage
+_global_logger = None
+
+
+def get_logger() -> "UnifiedLogger":
+    """Retrieve the global UnifiedLogger. If not initialized, fallback to auto-initialization."""
+    global _global_logger
+    if _global_logger is None:
+        try:
+            commit_hash = (
+                subprocess.check_output(
+                    ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.STDOUT
+                )
+                .decode("utf-8")
+                .strip()
+            )
+        except Exception:
+            commit_hash = "unknown"
+
+        timestamp = int(time.time())
+        generic_name = f"brittle_star_{commit_hash}_{timestamp}"
+
+        # Initialize generic fallback logger without WandB
+        _global_logger = UnifiedLogger(
+            run_name=generic_name,
+            config={"auto_initialized": True},
+            use_wandb=False,
+            _set_as_global=False,  # Prevent recursive call inside __init__
+        )
+        _global_logger.warning(f"UnifiedLogger auto-initialized with name: {generic_name}")
+
+    return _global_logger
 
 
 class UnifiedLogger:
@@ -33,6 +65,7 @@ class UnifiedLogger:
         base_dir: str = "runs",
         use_wandb: bool = True,
         save_code: bool = True,
+        _set_as_global: bool = True,
     ):
         """Initialize the unified logger.
 
@@ -44,6 +77,7 @@ class UnifiedLogger:
             base_dir: Base directory for local storage
             use_wandb: Whether to use WandB logging
             save_code: Whether to save code to WandB
+            _set_as_global: Internal flag to override the global singleton
         """
         self.run_name = run_name
         self.config = config
@@ -63,6 +97,28 @@ class UnifiedLogger:
 
         self.config_file = self.run_dir / "config.json"
 
+        # Setup standard Python logging mirror
+        self.text_log_file = self.run_dir / "run.log"
+        self._text_logger = logging.getLogger(f"UnifiedLogger_{self.run_name}")
+        self._text_logger.setLevel(logging.INFO)
+
+        # Avoid duplicate handlers if re-instantiated
+        if not self._text_logger.handlers:
+            fh = logging.FileHandler(self.text_log_file)
+            ch = logging.StreamHandler()
+
+            formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+            fh.setFormatter(formatter)
+            ch.setFormatter(formatter)
+
+            self._text_logger.addHandler(fh)
+            self._text_logger.addHandler(ch)
+
+        # Set as global singleton
+        global _global_logger
+        if _set_as_global:
+            _global_logger = self
+
         # Save config to disk
         self._save_config()
 
@@ -71,12 +127,28 @@ class UnifiedLogger:
             self._init_wandb(project_name, entity, save_code)
 
         # Initialize metrics storage
-        self.metrics_buffer: list[Dict[str, Any]] = []
+        self.metrics_buffer: List[Dict[str, Any]] = []
         self.step_counter = 0
 
-        logger.info(f"Initialized for run: {run_name}")
+        self.info(f"Initialized UnifiedLogger for run: {run_name}")
-        logger.info(f"Local storage: {self.run_dir.absolute()}")
+        self.info(f"Local storage: {self.run_dir.absolute()}")
-        logger.info(f"WandB logging: {self.wandb_available}")
+        self.info(f"WandB logging: {self.wandb_available}")
+
+    def info(self, msg: str, *args, **kwargs):
+        """Log an info message to stdout and disk."""
+        self._text_logger.info(msg, *args, **kwargs)
+
+    def warning(self, msg: str, *args, **kwargs):
+        """Log a warning message to stdout and disk."""
+        self._text_logger.warning(msg, *args, **kwargs)
+
+    def error(self, msg: str, *args, **kwargs):
+        """Log an error message to stdout and disk."""
+        self._text_logger.error(msg, *args, **kwargs)
+
+    def debug(self, msg: str, *args, **kwargs):
+        """Log a debug message to stdout and disk."""
+        self._text_logger.debug(msg, *args, **kwargs)
 
     def _init_wandb(self, project_name: str, entity: Optional[str], save_code: bool):
         """Initialize Weights & Biases logging."""
@@ -95,9 +167,9 @@ class UnifiedLogger:
         try:
             with open(self.config_file, "w") as f:
                 json.dump(self.config, f, indent=2)
-            logger.info(f"Config saved to {self.config_file}")
+            self.info(f"Config saved to {self.config_file}")
         except Exception as e:
-            logger.error(f"Error saving config: {e}")
+            self.error(f"Error saving config: {e}")
 
     def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True):
         """Log metrics to all backends.
@@ -126,7 +198,7 @@ class UnifiedLogger:
             try:
                 self.wandb_run.log(metrics, step=step, commit=commit)
             except Exception as e:
-                logger.warning(f"WandB logging failed: {e}")
+                self.warning(f"WandB logging failed: {e}")
 
         # Buffer for disk storage
         self.metrics_buffer.append(metrics_with_metadata)
@@ -143,7 +215,7 @@ class UnifiedLogger:
             for k, v in metrics.items()
             if k not in ["step", "timestamp"]
         )
-        logger.info(f"[Step {step}] {metric_str}")
+        self.info(f"[Step {step}] {metric_str}")
 
     def _flush_metrics(self):
         """Flush buffered metrics to disk."""
@@ -166,7 +238,7 @@ class UnifiedLogger:
                     f.write(json.dumps(serializable_metric) + "\n")
             self.metrics_buffer.clear()
         except Exception as e:
-            logger.error(f"Error flushing metrics: {e}")
+            self.error(f"Error flushing metrics: {e}")
 
     def save_checkpoint(
         self,
@@ -175,14 +247,7 @@ class UnifiedLogger:
         prefix: str = "checkpoint",
         metadata: Optional[Dict[str, Any]] = None,
     ):
-        """Save model checkpoint to disk and optionally to WandB.
+        """Save model checkpoint to disk and optionally to WandB."""
-
-        Args:
-            params: Model parameters (Flax params or any serializable object)
-            step: Current training step
-            prefix: Prefix for checkpoint filename
-            metadata: Additional metadata to save with checkpoint
-        """
         checkpoint_name = f"{prefix}_step_{step}.flax"
         checkpoint_path = self.checkpoints_dir / checkpoint_name
 
@@ -197,7 +262,7 @@ class UnifiedLogger:
                 with open(metadata_path, "w") as f:
                     json.dump(metadata, f, indent=2)
 
-            logger.info(f"Checkpoint saved: {checkpoint_path}")
+            self.info(f"Checkpoint saved: {checkpoint_path}")
 
             # Log to WandB as artifact
             if self.wandb_available:
@@ -213,20 +278,15 @@ class UnifiedLogger:
                     if metadata:
                         artifact.add_file(str(metadata_path))
                     self.wandb_run.log_artifact(artifact)
-                    logger.info("Checkpoint uploaded to WandB")
+                    self.info("Checkpoint uploaded to WandB")
                 except Exception as e:
-                    logger.warning(f"Could not upload checkpoint to WandB: {e}")
+                    self.warning(f"Could not upload checkpoint to WandB: {e}")
 
         except Exception as e:
-            logger.error(f"Error saving checkpoint: {e}")
+            self.error(f"Error saving checkpoint: {e}")
 
     def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None):
-        """Save the final trained model.
+        """Save the final trained model."""
-
-        Args:
-            params: Model parameters
-            metadata: Additional metadata about the final model
-        """
         final_model_path = self.run_dir / "final_model.flax"
 
         try:
@@ -238,7 +298,7 @@ class UnifiedLogger:
                 with open(metadata_path, "w") as f:
                     json.dump(metadata, f, indent=2)
 
-            logger.info(f"Final model saved: {final_model_path}")
+            self.info(f"Final model saved: {final_model_path}")
 
             # Log to WandB
             if self.wandb_available:
@@ -255,17 +315,17 @@ class UnifiedLogger:
                         artifact.add_file(str(metadata_path))
                     self.wandb_run.log_artifact(artifact)
                 except Exception as e:
-                    logger.warning(f"Could not upload final model to WandB: {e}")
+                    self.warning(f"Could not upload final model to WandB: {e}")
 
         except Exception as e:
-            logger.error(f"Error saving final model: {e}")
+            self.error(f"Error saving final model: {e}")
 
     def finish(self):
         """Finalize logging and cleanup."""
         # Flush remaining metrics
         self._flush_metrics()
 
-        logger.info(f"Run complete. Results saved to: {self.run_dir.absolute()}")
+        self.info(f"Run complete. Results saved to: {self.run_dir.absolute()}")
 
         # Finish WandB run
         if self.wandb_available:

src/main.py

@@ -1,9 +1,6 @@
-import logging
-
 import jax
-
+from experiment_logger import get_logger
-logger = logging.getLogger(__name__)
 
 if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
+    logger = get_logger()
     logger.info(f"JAX devices: {jax.devices()}")

src/train.py

@@ -1,4 +1,3 @@
-import logging
 import random
 import time
 from dataclasses import asdict
@@ -12,7 +11,6 @@ import numpy as np
 import optax
 import torch
 import tqdm
-import tyro
 from flax.training.train_state import TrainState
 from torch.utils.tensorboard import SummaryWriter
 
@@ -20,10 +18,8 @@ from brittle_star_project.dataclasses import PPOArgs
 from brittle_star_project.dataclasses.EpisodeStatistics import EpisodeStatistics
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 from brittle_star_project.rl import Actor, AgentParams, Critic, Network, Storage
-from experiment_logger import UnifiedLogger
+from experiment_logger import UnifiedLogger, get_logger
-from experiment_logger.config_utils import merge_config_with_cli, print_config
+from experiment_logger.config_utils import merge_config_with_cli
-
-log = logging.getLogger(__name__)
 
 
 def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
@@ -44,7 +40,7 @@ def train(args: PPOArgs):
     args.minibatch_size = args.batch_size // args.num_minibatches
     args.num_iterations = args.total_timesteps // args.batch_size
     run_name = f"{args.exp_name}__seed_{args.seed}__{int(time.time())}"
-    log.info(f"Run name: {run_name}")
+    get_logger().info(f"Run name: {run_name}")
 
     # Initialize unified logger (replaces wandb.init and tensorboard writer)
     logger = UnifiedLogger(
@@ -71,9 +67,9 @@ def train(args: PPOArgs):
     torch.backends.cudnn.deterministic = args.torch_deterministic
     device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
     device = "cpu"  # Force CPU for JAX
-    log.info(f"Device: {device}")
+    logger.info(f"Device: {device}")
 
-    log.info("Creating environment...")
+    logger.info("Creating environment...")
     env = make_env(num_envs=args.num_envs)()
 
     episode_stats = EpisodeStatistics(
@@ -115,7 +111,7 @@ def train(args: PPOArgs):
         frac = 1.0 - (count // (args.num_minibatches * args.update_epochs)) / args.num_iterations
         return args.learning_rate * frac
 
-    log.info("Initializing models...")
+    logger.info("Initializing models...")
     network = Network()
     actor = Actor(action_dim=env.single_action_space.shape[0])  # continuous actions for MJX
     critic = Critic()
@@ -264,7 +260,7 @@ def train(args: PPOArgs):
     start_time = time.time()
 
     # Reset once to get initial state
-    log.info("Resetting environment...")
+    logger.info("Resetting environment...")
     next_env_state = env.reset(seed=args.seed)
     next_obs = convert_obs_dict_to_array(next_env_state.observations)
     next_done = jnp.zeros(args.num_envs, dtype=jnp.bool_)
@@ -306,7 +302,7 @@ def train(args: PPOArgs):
         max_steps=args.num_steps,
     )
 
-    log.info("Starting training...")
+    logger.info("Starting training...")
     iters_bar = tqdm.tqdm(range(1, args.num_iterations + 1))
     for iteration in iters_bar:
         iteration_time_start = time.time()
@@ -406,7 +402,7 @@ def train(args: PPOArgs):
                     ]
                 )
             )
-        log.info(f"Legacy model saved to {model_path}")
+        logger.info(f"Legacy model saved to {model_path}")
 
     # Finalize logging
     logger.finish()
@@ -415,15 +411,12 @@ def train(args: PPOArgs):
 
 
 def main() -> None:
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    )
-    
     # Enhanced argument parsing with YAML config support
     args = merge_config_with_cli(PPOArgs)
 
     # Print final configuration
+    from experiment_logger.config_utils import print_config
+
     print_config(args, "Final Training Configuration")
 
     train(args)