tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

refactor(log): improved logging workflow

Browse source
This commit is contained in:
Tibo De Peuter 2026-04-01 16:01:12 +00:00
parent 7e7c5bf27c
commit ff90101377
7 changed files with 162 additions and 90 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
View file

@ -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)

8
src/brittle_star_project/environment/factory.py
View file

@ -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

11
src/experiment_logger/__init__.py
View file

@ -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"

57
src/experiment_logger/config_utils.py
View file

@ -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}")
log.info(f" {key}: {value}")

128
src/experiment_logger/unified_logger.py
View file

@ -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:

7
src/main.py
View file

@ -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()}")

33
src/train.py
View file

@ -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)