diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 1ea07c7..c7cb3c6 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/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) diff --git a/src/brittle_star_project/environment/factory.py b/src/brittle_star_project/environment/factory.py index 1cf94f2..1a891ea 100644 --- a/src/brittle_star_project/environment/factory.py +++ b/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 diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index afbb13b..93b1de3 100644 --- a/src/experiment_logger/__init__.py +++ b/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" diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index b296266..7a92c0e 100644 --- a/src/experiment_logger/config_utils.py +++ b/src/experiment_logger/config_utils.py @@ -1,28 +1,29 @@ """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]: """Load configuration from YAML file.""" 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: return {} - + log.info(f"Loaded configuration from: {config_path}") return config @@ -30,10 +31,10 @@ def load_yaml_config(config_path: str) -> Dict[str, Any]: 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}") @@ -41,10 +42,10 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: """Create dataclass instance from dictionary, handling type conversions.""" if not is_dataclass(cls): raise ValueError(f"{cls} is not a dataclass") - + # Get field names and types field_map = {f.name: f for f in fields(cls)} - + # Filter config to only include valid fields filtered_config = {} for key, value in config_dict.items(): @@ -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): - filtered_config[key] = value.lower() in ('true', '1', 'yes', 'on') + if field.type is bool and isinstance(value, str): + 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: @@ -69,23 +70,23 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: filtered_config[key] = value else: log.warning(f"Unknown configuration parameter: {key}") - + return cls(**filtered_config) def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T: """Merge YAML config with CLI arguments, with CLI taking precedence. - + Args: config_class: Dataclass type to create config_file: Path to YAML config file (optional) - + Returns: Instance of config_class with merged configuration """ # Parse CLI args first to get the default/CLI values import tyro - + # Check if --config is in sys.argv and extract it extracted_config_file = config_file if "--config" in sys.argv: @@ -95,7 +96,7 @@ def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T: # Remove from sys.argv so tyro doesn't see it sys.argv.pop(config_idx) # Remove --config sys.argv.pop(config_idx) # Remove config file path - + # Load YAML config if available yaml_config = {} if extracted_config_file and os.path.exists(extracted_config_file): @@ -103,24 +104,24 @@ def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T: log.info(f"Merging YAML config from {extracted_config_file} with CLI args") elif extracted_config_file: log.warning(f"Config file not found: {extracted_config_file}, using CLI args only") - + # Create default instance to know what the defaults are default_instance = config_class() default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)} - + # Parse CLI args cli_instance = tyro.cli(config_class) cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)} - + # Merge configs: YAML as base, CLI overrides non-default values final_config = {} - + for field in fields(config_class): field_name = field.name - default_value = default_dict[field_name] + default_value = default_dict[field_name] yaml_value = yaml_config.get(field_name, default_value) cli_value = cli_dict[field_name] - + # Use CLI value if it's different from default, otherwise use YAML value if cli_value != default_value: final_config[field_name] = cli_value @@ -130,7 +131,7 @@ def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T: final_config[field_name] = yaml_value if yaml_value != default_value: log.info(f"YAML config: {field_name}={yaml_value}") - + return config_class(**final_config) @@ -143,4 +144,4 @@ def print_config(config: Any, title: str = "Configuration"): log.info(f" {field.name}: {value}") else: for key, value in vars(config).items(): - log.info(f" {key}: {value}") \ No newline at end of file + log.info(f" {key}: {value}") diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 1501bf5..c391f52 100644 --- a/src/experiment_logger/unified_logger.py +++ b/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}") - logger.info(f"Local storage: {self.run_dir.absolute()}") - logger.info(f"WandB logging: {self.wandb_available}") + self.info(f"Initialized UnifiedLogger for run: {run_name}") + self.info(f"Local storage: {self.run_dir.absolute()}") + 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. - - 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 - """ + """Save model checkpoint to disk and optionally to WandB.""" 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. - - Args: - params: Model parameters - metadata: Additional metadata about the final model - """ + """Save the final trained 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: diff --git a/src/main.py b/src/main.py index ac40465..61712ce 100644 --- a/src/main.py +++ b/src/main.py @@ -1,9 +1,6 @@ -import logging - import jax - -logger = logging.getLogger(__name__) +from experiment_logger import get_logger if __name__ == "__main__": - logging.basicConfig(level=logging.INFO) + logger = get_logger() logger.info(f"JAX devices: {jax.devices()}") diff --git a/src/train.py b/src/train.py index ad31b56..e40e751 100644 --- a/src/train.py +++ b/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.config_utils import merge_config_with_cli, print_config - -log = logging.getLogger(__name__) +from experiment_logger import UnifiedLogger, get_logger +from experiment_logger.config_utils import merge_config_with_cli 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,17 +411,14 @@ 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)