From c27f5fcdf2188ded3bb88cad6d355af3ce6ba9ef Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:32:38 +0000 Subject: [PATCH 01/31] fix(lint): remove deprecated ruff rule PLR1708 --- ruff.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ruff.toml b/ruff.toml index b19b931..16e880c 100644 --- a/ruff.toml +++ b/ruff.toml @@ -349,7 +349,7 @@ extend-ignore = [ # "PLR1705", # no-else-return # "PLR1706", # consider-using-ternary # "PLR1707", # trailing-comma-tuple - "PLR1708", # stop-iteration-return + # "PLR1708", # stop-iteration-return (deprecated) # "PLR1709", # simplify-boolean-expression # "PLR1710", # inconsistent-return-statements "PLR1711", # useless-return From 3ce107a5606caef50f1345ca57a8c4dedc4281f9 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:49:04 +0000 Subject: [PATCH 02/31] feat(experiment-logger): add standalone logging framework Create reusable experiment logging package with: - UnifiedLogger for multi-backend logging (WandB, disk, stdout) - Automatic checkpoint and model saving with metadata - WandB artifact upload support - Graceful degradation when WandB unavailable - Comprehensive API documentation This is a standalone, project-agnostic library that can be reused across different ML projects. --- src/experiment_logger/README.md | 145 +++++++++++++ src/experiment_logger/__init__.py | 11 + src/experiment_logger/unified_logger.py | 270 ++++++++++++++++++++++++ src/experiment_logger/wandb_utils.py | 69 ++++++ 4 files changed, 495 insertions(+) create mode 100644 src/experiment_logger/README.md create mode 100644 src/experiment_logger/__init__.py create mode 100644 src/experiment_logger/unified_logger.py create mode 100644 src/experiment_logger/wandb_utils.py diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md new file mode 100644 index 0000000..2beac1a --- /dev/null +++ b/src/experiment_logger/README.md @@ -0,0 +1,145 @@ +# Experiment Logger + +A lightweight, standalone logging framework for machine learning experiments with multi-backend support. + +## Features + +- **Multi-backend logging**: Simultaneously log to WandB, local disk (JSON), and stdout +- **Data preservation**: All metrics saved locally, even if WandB is unavailable +- **Checkpoint management**: Save model checkpoints with metadata +- **WandB integration**: Optional artifact upload for model versioning +- **Graceful degradation**: Works without WandB installed +- **Simple API**: Minimal configuration required + +## Installation + +This package is included in the project. To use it in your code: + +```python +from experiment_logger import UnifiedLogger +``` + +## Quick Start + +```python +from experiment_logger import UnifiedLogger + +# Initialize logger +logger = UnifiedLogger( + run_name="my_experiment", + config={"learning_rate": 0.001, "batch_size": 32}, + project_name="MyProject", + entity="my-wandb-username", # Optional + use_wandb=True, # Set to False to disable WandB +) + +# Log metrics +for step in range(100): + logger.log({ + "loss": 1.0 / (step + 1), + "accuracy": step * 0.01, + }, step=step) + +# Save checkpoint +logger.save_checkpoint( + params=model_params, + step=100, + metadata={"epoch": 1, "val_acc": 0.95}, +) + +# Save final model +logger.save_final_model( + params=final_params, + metadata={"final_accuracy": 0.98}, +) + +# Finalize (flushes remaining metrics) +logger.finish() +``` + +## Context Manager + +Use as a context manager for automatic cleanup: + +```python +with UnifiedLogger(run_name="my_exp", config={}) as logger: + logger.log({"metric": 1.0}) + # Automatically calls finish() on exit +``` + +## Configuration + +### Constructor Parameters + +- `run_name` (str): Unique name for this run +- `config` (dict): Configuration dictionary with hyperparameters +- `project_name` (str): WandB project name (default: "PPO-Modularity") +- `entity` (str, optional): WandB entity (team/user name) +- `base_dir` (str): Base directory for local storage (default: "runs") +- `use_wandb` (bool): Enable WandB logging (default: True) +- `save_code` (bool): Save code to WandB (default: True) + +### Directory Structure + +``` +runs/ +└── my_experiment/ + ├── config.json # Saved configuration + ├── metrics/ + │ └── metrics.jsonl # Line-delimited JSON metrics + ├── checkpoints/ + │ ├── checkpoint_step_100.flax + │ └── checkpoint_step_100_metadata.json + └── final_model.flax +``` + +## API Reference + +### `log(metrics, step=None, commit=True)` + +Log metrics to all backends. + +**Parameters:** +- `metrics` (dict): Dictionary of metric name -> value +- `step` (int, optional): Global step counter (auto-incremented if None) +- `commit` (bool): Whether to commit to WandB immediately + +### `save_checkpoint(params, step, prefix="checkpoint", metadata=None)` + +Save model checkpoint to disk and optionally to WandB. + +**Parameters:** +- `params`: Model parameters (Flax params or any serializable object) +- `step` (int): Current training step +- `prefix` (str): Prefix for checkpoint filename +- `metadata` (dict, optional): Additional metadata to save + +### `save_final_model(params, metadata=None)` + +Save the final trained model. + +**Parameters:** +- `params`: Model parameters +- `metadata` (dict, optional): Metadata about the final model + +### `finish()` + +Finalize logging and cleanup. Flushes remaining metrics to disk. + +## Usage in Projects + +This logger is designed to be: +- **Project-agnostic**: Use in any ML project, not just this one +- **Framework-agnostic**: Works with JAX, PyTorch, TensorFlow, etc. +- **Minimal dependencies**: Only requires `wandb` (optional), `flax` (for serialization), and `numpy` + +## Design Philosophy + +1. **Never lose data**: All metrics saved locally, regardless of WandB availability +2. **Simple API**: Minimal boilerplate, easy to integrate +3. **Fail gracefully**: Missing WandB shouldn't break experiments +4. **Reproducibility**: Save full configuration with every run + +## License + +Part of the 2026SEL3-project-BrittleStar repository. diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py new file mode 100644 index 0000000..7ac168d --- /dev/null +++ b/src/experiment_logger/__init__.py @@ -0,0 +1,11 @@ +"""Unified logging framework for machine learning experiments. + +This package provides a unified interface for logging to multiple backends +(WandB, disk, stdout) simultaneously, ensuring no data loss. +""" + +from experiment_logger.unified_logger import UnifiedLogger +from experiment_logger.wandb_utils import finish_wandb, init_wandb + +__all__ = ["UnifiedLogger", "init_wandb", "finish_wandb"] +__version__ = "0.1.0" diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py new file mode 100644 index 0000000..87e3006 --- /dev/null +++ b/src/experiment_logger/unified_logger.py @@ -0,0 +1,270 @@ +"""Unified logger that writes to multiple backends simultaneously. + +This logger ensures all experimental data is preserved by writing to: +1. Weights & Biases (when available) +2. Local disk (JSON files, model checkpoints) +3. stdout (for real-time monitoring) +""" + +import json +import logging +import time +from pathlib import Path +from typing import Any, Dict, Optional + +import flax +import numpy as np + +from experiment_logger.wandb_utils import finish_wandb, init_wandb + +logger = logging.getLogger(__name__) + + +class UnifiedLogger: + """Unified logger for scientific experiments with redundant backup.""" + + def __init__( + self, + run_name: str, + config: Dict[str, Any], + project_name: str = "PPO-Modularity", + entity: Optional[str] = None, + base_dir: str = "runs", + use_wandb: bool = True, + save_code: bool = True, + ): + """Initialize the unified logger. + + Args: + run_name: Unique name for this run + config: Configuration dictionary with hyperparameters + project_name: WandB project name + entity: WandB entity (team/user name) + 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.wandb_available = False + self.wandb_run = None + + # Setup local storage + self.run_dir = Path(base_dir) / run_name + self.run_dir.mkdir(parents=True, exist_ok=True) + + self.checkpoints_dir = self.run_dir / "checkpoints" + self.checkpoints_dir.mkdir(exist_ok=True) + + self.metrics_dir = self.run_dir / "metrics" + self.metrics_dir.mkdir(exist_ok=True) + + self.config_file = self.run_dir / "config.json" + + # Save config to disk + self._save_config() + + # Initialize WandB if requested + if self.use_wandb: + self._init_wandb(project_name, entity, save_code) + + # Initialize metrics storage + 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}") + + def _init_wandb(self, project_name: str, entity: Optional[str], save_code: bool): + """Initialize Weights & Biases logging.""" + self.wandb_run = init_wandb( + project=project_name, + entity=entity, + name=self.run_name, + config=self.config, + save_code=save_code, + resume="allow", + ) + self.wandb_available = self.wandb_run is not None + + def _save_config(self): + """Save configuration to disk.""" + 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}") + except Exception as e: + logger.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. + + Args: + metrics: Dictionary of metric name -> value + step: Global step counter (auto-incremented if None) + commit: Whether to commit to WandB immediately + """ + if step is None: + step = self.step_counter + self.step_counter += 1 + + # Add timestamp + metrics_with_metadata = { + "step": step, + "timestamp": time.time(), + **metrics, + } + + # Log to stdout + self._log_to_stdout(metrics_with_metadata) + + # Log to WandB + if self.wandb_available: + try: + self.wandb_run.log(metrics, step=step, commit=commit) + except Exception as e: + logger.warning(f"WandB logging failed: {e}") + + # Buffer for disk storage + self.metrics_buffer.append(metrics_with_metadata) + + # Periodically flush to disk + if len(self.metrics_buffer) >= 100: + self._flush_metrics() + + def _log_to_stdout(self, metrics: Dict[str, Any]): + """Log metrics to stdout for real-time monitoring.""" + step = metrics.get("step", "?") + metric_str = ", ".join( + f"{k}={v:.6f}" if isinstance(v, (float, np.floating)) else f"{k}={v}" + for k, v in metrics.items() + if k not in ["step", "timestamp"] + ) + logger.info(f"[Step {step}] {metric_str}") + + def _flush_metrics(self): + """Flush buffered metrics to disk.""" + if not self.metrics_buffer: + return + + try: + metrics_file = self.metrics_dir / "metrics.jsonl" + with open(metrics_file, "a") as f: + for metric in self.metrics_buffer: + f.write(json.dumps(metric) + "\n") + self.metrics_buffer.clear() + except Exception as e: + logger.error(f"Error flushing metrics: {e}") + + def save_checkpoint( + self, + params: Any, + step: int, + 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 + """ + checkpoint_name = f"{prefix}_step_{step}.flax" + checkpoint_path = self.checkpoints_dir / checkpoint_name + + try: + # Save to disk using Flax serialization + with open(checkpoint_path, "wb") as f: + f.write(flax.serialization.to_bytes(params)) + + # Save metadata if provided + if metadata: + metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.json" + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + + logger.info(f"Checkpoint saved: {checkpoint_path}") + + # Log to WandB as artifact + if self.wandb_available: + try: + import wandb + + artifact = wandb.Artifact( + name=f"{self.run_name}_{prefix}", + type="model", + metadata=metadata or {}, + ) + artifact.add_file(str(checkpoint_path)) + if metadata: + artifact.add_file(str(metadata_path)) + self.wandb_run.log_artifact(artifact) + logger.info("Checkpoint uploaded to WandB") + except Exception as e: + logger.warning(f"Could not upload checkpoint to WandB: {e}") + + except Exception as e: + logger.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 + """ + final_model_path = self.run_dir / "final_model.flax" + + try: + with open(final_model_path, "wb") as f: + f.write(flax.serialization.to_bytes(params)) + + if metadata: + metadata_path = self.run_dir / "final_model_metadata.json" + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + + logger.info(f"Final model saved: {final_model_path}") + + # Log to WandB + if self.wandb_available: + try: + import wandb + + artifact = wandb.Artifact( + name=f"{self.run_name}_final_model", + type="model", + metadata=metadata or {}, + ) + artifact.add_file(str(final_model_path)) + if metadata: + 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}") + + except Exception as e: + logger.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()}") + + # Finish WandB run + if self.wandb_available: + finish_wandb() + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.finish() diff --git a/src/experiment_logger/wandb_utils.py b/src/experiment_logger/wandb_utils.py new file mode 100644 index 0000000..5fd8837 --- /dev/null +++ b/src/experiment_logger/wandb_utils.py @@ -0,0 +1,69 @@ +"""Centralized WandB initialization utilities.""" + +import logging +from typing import Any, Dict, Optional + +logger = logging.getLogger(__name__) + + +def init_wandb( + project: str, + config: Dict[str, Any], + name: Optional[str] = None, + entity: Optional[str] = None, + sync_tensorboard: bool = False, + save_code: bool = True, + resume: str = "allow", + **kwargs, +): + """Initialize WandB with standardized settings. + + This function provides a centralized way to initialize WandB across different + scripts, ensuring consistent configuration and error handling. + + Args: + project: WandB project name + config: Configuration dictionary to log + name: Run name (auto-generated if None) + entity: WandB entity (team/user name) + sync_tensorboard: Whether to sync tensorboard logs + save_code: Whether to save code snapshots + resume: Resume strategy ("allow", "must", "never", "auto") + **kwargs: Additional arguments to pass to wandb.init() + + Returns: + wandb.Run object if successful, None otherwise + """ + try: + import wandb + + run = wandb.init( + project=project, + entity=entity, + name=name, + config=config, + sync_tensorboard=sync_tensorboard, + save_code=save_code, + resume=resume, + **kwargs, + ) + logger.info(f"WandB initialized successfully for project '{project}', run '{run.name}'") + return run + except ImportError: + logger.warning("WandB not installed. Skipping WandB initialization.") + return None + except Exception as e: + logger.error(f"Failed to initialize WandB: {e}") + return None + + +def finish_wandb(): + """Safely finish the current WandB run.""" + try: + import wandb + + if wandb.run is not None: + wandb.finish() + logger.info("WandB run finished successfully") + except Exception as e: + logger.warning(f"Error finishing WandB run: {e}") From bb0bb94f600392d49b8f46009cc4a023c406fcb4 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:49:12 +0000 Subject: [PATCH 03/31] feat(config): add checkpoint frequency parameter Add checkpoint_frequency to PPOArgs to enable periodic checkpoint saving during training. Default: save every 100 iterations. --- src/brittle_star_project/dataclasses/PPOArgs.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 4fc3f2e..b27265f 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/src/brittle_star_project/dataclasses/PPOArgs.py @@ -26,7 +26,7 @@ class PPOArgs: wandb_project_name: str = "PPO-Modularity" # the entity (team) of wandb's project - wandb_entity: str | None = None + wandb_entity: str | None = "tdpeuter-ghent-university" # whether to capture videos of the agent performances (check out `videos` folder) capture_video: bool = False @@ -34,6 +34,9 @@ class PPOArgs: # whether to save model into the `runs/{run_name}` folder save_model: bool = True + # checkpoint frequency (in iterations, 0 = no intermediate checkpoints) + checkpoint_frequency: int = 100 + # whether to upload the saved model to huggingface upload_model: bool = False From b70cd1c27b9ae538cb604790c0565ef381a60a1d Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:49:21 +0000 Subject: [PATCH 04/31] feat(train): integrate experiment_logger and replace print statements - Replace direct WandB calls with experiment_logger.UnifiedLogger - Replace all print() calls with proper logging framework - Add automatic checkpoint saving every N iterations - Configure root logger with proper format and level - Maintain backward compatibility with TensorBoard writer - Save final model with metadata using unified logger --- src/train.py | 125 +++++++++++++++++++++++++++++++++++---------------- 1 file changed, 87 insertions(+), 38 deletions(-) diff --git a/src/train.py b/src/train.py index eb64bbe..00e8c60 100644 --- a/src/train.py +++ b/src/train.py @@ -1,3 +1,4 @@ +import logging import random import time from dataclasses import asdict @@ -18,7 +19,10 @@ from torch.utils.tensorboard import SummaryWriter 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 Network, Actor, Critic, AgentParams, Storage +from brittle_star_project.rl import Actor, AgentParams, Critic, Network, Storage +from experiment_logger import UnifiedLogger + +log = logging.getLogger(__name__) def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: @@ -39,20 +43,19 @@ 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())}" - print(f"running name: {run_name}") + log.info(f"Run name: {run_name}") - if args.track: - import wandb - - wandb.init( - project=args.wandb_project_name, - entity=args.wandb_entity, - sync_tensorboard=True, - config=vars(args), - name=run_name, - save_code=True, - ) + # Initialize unified logger (replaces wandb.init and tensorboard writer) + logger = UnifiedLogger( + run_name=run_name, + config=vars(args), + project_name=args.wandb_project_name, + entity=args.wandb_entity, + use_wandb=args.track, + save_code=True, + ) + # Keep TensorBoard writer for backward compatibility writer = SummaryWriter(f"runs/{run_name}") writer.add_text( "hyperparameters", @@ -66,9 +69,10 @@ 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") - print(f"Running on device: {device}") + device = "cpu" # Force CPU for JAX + log.info(f"Device: {device}") - print("Creating the environment...") + log.info("Creating environment...") env = make_env(num_envs=args.num_envs)() episode_stats = EpisodeStatistics( @@ -110,7 +114,7 @@ def train(args: PPOArgs): frac = 1.0 - (count // (args.num_minibatches * args.update_epochs)) / args.num_iterations return args.learning_rate * frac - print("Initializing the models...") + log.info("Initializing models...") network = Network() actor = Actor(action_dim=env.single_action_space.shape[0]) # continuous actions for MJX critic = Critic() @@ -259,7 +263,7 @@ def train(args: PPOArgs): start_time = time.time() # Reset once to get initial state - print("Resetting the environment...") + log.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_) @@ -301,9 +305,9 @@ def train(args: PPOArgs): max_steps=args.num_steps, ) - print("Starting training...") + log.info("Starting training...") iters_bar = tqdm.tqdm(range(1, args.num_iterations + 1)) - for _ in iters_bar: + for iteration in iters_bar: iteration_time_start = time.time() agent_state, episode_stats, next_obs, next_done, storage, key, next_env_state = rollout( @@ -317,37 +321,76 @@ def train(args: PPOArgs): ) avg_episodic_return = np.mean(jax.device_get(episode_stats.returned_episode_returns)) + avg_episodic_length = np.mean(jax.device_get(episode_stats.returned_episode_lengths)) + learning_rate = agent_state.opt_state[1].hyperparams["learning_rate"].item() + sps = int(global_step / (time.time() - start_time)) + sps_update = int(args.num_envs * args.num_steps / (time.time() - iteration_time_start)) + iters_bar.set_postfix_str( f"global_step={global_step}, avg_episodic_return={avg_episodic_return}" ) + # Log to unified logger + logger.log( + { + "charts/avg_episodic_return": avg_episodic_return, + "charts/avg_episodic_length": avg_episodic_length, + "charts/learning_rate": learning_rate, + "charts/SPS": sps, + "charts/SPS_update": sps_update, + "losses/value_loss": v_loss[-1, -1].item(), + "losses/policy_loss": pg_loss[-1, -1].item(), + "losses/entropy": entropy_loss[-1, -1].item(), + "losses/approx_kl": approx_kl[-1, -1].item(), + "losses/loss": loss[-1, -1].item(), + }, + step=global_step, + ) + + # Also log to TensorBoard for backward compatibility writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step) - writer.add_scalar( - "charts/avg_episodic_length", - np.mean(jax.device_get(episode_stats.returned_episode_lengths)), - global_step, - ) - writer.add_scalar( - "charts/learning_rate", - agent_state.opt_state[1].hyperparams["learning_rate"].item(), - global_step, - ) + writer.add_scalar("charts/avg_episodic_length", avg_episodic_length, global_step) + writer.add_scalar("charts/learning_rate", learning_rate, global_step) writer.add_scalar("losses/value_loss", v_loss[-1, -1].item(), global_step) writer.add_scalar("losses/policy_loss", pg_loss[-1, -1].item(), global_step) writer.add_scalar("losses/entropy", entropy_loss[-1, -1].item(), global_step) writer.add_scalar("losses/approx_kl", approx_kl[-1, -1].item(), global_step) writer.add_scalar("losses/loss", loss[-1, -1].item(), global_step) + writer.add_scalar("charts/SPS", sps, global_step) + writer.add_scalar("charts/SPS_update", sps_update, global_step) - # iters_bar.set_postfix_str(f"SPS: {int(global_step / (time.time() - start_time))}") - - writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step) - writer.add_scalar( - "charts/SPS_update", - int(args.num_envs * args.num_steps / (time.time() - iteration_time_start)), - global_step, - ) + # Save periodic checkpoints + if args.checkpoint_frequency > 0 and iteration % args.checkpoint_frequency == 0: + logger.save_checkpoint( + params={ + "network_params": agent_state.params["network_params"], + "actor_params": agent_state.params["actor_params"], + "critic_params": agent_state.params["critic_params"], + }, + step=global_step, + metadata={ + "iteration": iteration, + "avg_episodic_return": float(avg_episodic_return), + "avg_episodic_length": float(avg_episodic_length), + }, + ) if args.save_model: + # Save using unified logger (better organization and WandB integration) + logger.save_final_model( + params={ + "network_params": agent_state.params["network_params"], + "actor_params": agent_state.params["actor_params"], + "critic_params": agent_state.params["critic_params"], + }, + metadata={ + "global_step": global_step, + "avg_episodic_return": float(avg_episodic_return), + "config": vars(args), + }, + ) + + # Also save in old format for backward compatibility model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model" with open(model_path, "wb") as f: f.write( @@ -362,13 +405,19 @@ def train(args: PPOArgs): ] ) ) - print(f"model saved to {model_path}") + log.info(f"Legacy model saved to {model_path}") + # Finalize logging + logger.finish() env.close() writer.close() def main() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) args = tyro.cli(PPOArgs) train(args) From 4151d2307e07a8b50ce55a13af99711b78753806 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:51:17 +0000 Subject: [PATCH 05/31] refactor(main): replace print with logging --- src/main.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/main.py b/src/main.py index ef7c36e..ac40465 100644 --- a/src/main.py +++ b/src/main.py @@ -1,4 +1,9 @@ +import logging + import jax +logger = logging.getLogger(__name__) + if __name__ == "__main__": - print(jax.devices()) + logging.basicConfig(level=logging.INFO) + logger.info(f"JAX devices: {jax.devices()}") From f31436bccd14fdb7866f6d7a826167fcb23257a5 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:53:15 +0000 Subject: [PATCH 06/31] feat(config): add YAML configuration system - Add default_ppo.yaml template for training configurations - Create configs/README.md with usage documentation - Enable per-researcher configuration without code changes - Support YAML config files with CLI parameter overrides - Document how to set personal WandB credentials safely This allows researchers to maintain personal configs without committing credentials to the repository. --- configs/README.md | 48 ++++++++++++++++++++++++++++++++++++++++ configs/default_ppo.yaml | 48 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+) create mode 100644 configs/README.md create mode 100644 configs/default_ppo.yaml diff --git a/configs/README.md b/configs/README.md new file mode 100644 index 0000000..d06fbb2 --- /dev/null +++ b/configs/README.md @@ -0,0 +1,48 @@ +# Configuration Files + +This directory contains configuration files for training experiments. + +## Usage + +Configuration files use YAML format and allow you to specify all training parameters in one place. + +### Quick Start + +Copy the default configuration template: +```bash +cp configs/default_ppo.yaml configs/my_experiment.yaml +``` + +Edit `my_experiment.yaml` to customize your experiment settings, particularly: +- `wandb_entity`: Your WandB username or team name +- `track`: Set to `true` to enable WandB logging +- Training hyperparameters as needed + +Run training with your config: +```bash +python src/train.py --config configs/my_experiment.yaml +``` + +### Override Parameters + +You can override any parameter from the command line: +```bash +python src/train.py --config configs/my_experiment.yaml --learning-rate 0.001 --num-envs 32 +``` + +### Configuration for Different Users + +Each researcher should create their own config file with their WandB settings: +```yaml +# configs/researcher_name.yaml +exp_name: "researcher_name_experiment" +track: true +wandb_project_name: "PPO-Modularity" +wandb_entity: "your-wandb-username" # Change this! +``` + +This approach allows everyone to use the codebase without modifying source files. + +## Available Configurations + +- `default_ppo.yaml` - Default PPO training configuration template diff --git a/configs/default_ppo.yaml b/configs/default_ppo.yaml new file mode 100644 index 0000000..5775a84 --- /dev/null +++ b/configs/default_ppo.yaml @@ -0,0 +1,48 @@ +# PPO Training Configuration Template +# +# This file provides an example configuration for PPO training. +# Copy this file and modify it for your specific experiments. +# +# Usage: +# python src/train.py --config-path configs/my_config.yaml +# Or override specific parameters: +# python src/train.py --learning-rate 0.001 --num-envs 32 + +# Experiment settings +exp_name: "brittle_star_ppo" +seed: 1 + +# Tracking settings +track: false # Set to true to enable WandB logging +wandb_project_name: "PPO-Modularity" +wandb_entity: null # Set to your WandB username or team name + +# Model saving +save_model: true +checkpoint_frequency: 100 # Save checkpoint every N iterations (0 = no checkpoints) + +# Environment settings +num_envs: 16 + +# Training hyperparameters +total_timesteps: 10000000 +learning_rate: 0.00025 +num_steps: 128 +anneal_lr: true + +# PPO specific +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 4 +update_epochs: 4 +norm_adv: true +clip_coef: 0.1 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null + +# Hardware +cuda: true +torch_deterministic: true From 7fe217de0fe1f46241f4798d2cb2ec587e5e56c3 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:53:45 +0000 Subject: [PATCH 07/31] docs: update documentation for logging and configuration - Add Quick Start guide with installation and configuration steps - Document unified logging framework features - Explain configuration management for multiple researchers - Add project structure overview showing experiment_logger - Update training examples with new patterns - Add logging best practices to CONTRIBUTING.md --- README.md | 74 +++++++++++++++++++++++++++++++++++++++++--- docs/CONTRIBUTING.md | 9 ++++++ 2 files changed, 78 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 47b0552..74b001e 100644 --- a/README.md +++ b/README.md @@ -1,17 +1,81 @@ # Brittle Star -## Usage +Reinforcement learning research on brittle star locomotion using PPO. -### UV +## Quick Start -To set up the UV module, you can run the following command: +### Installation + +Set up the environment using UV: ```bash uv sync --frozen ``` -example command: +### Configuration + +1. **Copy the default configuration:** + ```bash + cp configs/default_ppo.yaml configs/my_experiment.yaml + ``` + +2. **Edit `configs/my_experiment.yaml`** to set your WandB credentials: + ```yaml + track: true # Enable WandB logging + wandb_entity: "your-wandb-username" # Replace with your username/team + wandb_project_name: "PPO-Modularity" + ``` + +3. **(Optional) Login to WandB:** + ```bash + uv run wandb login + ``` + +### Training + +Run training with your configuration: ```bash -uv run src/train.py --model_name my_model --epochs 50 --batch_size 32 +uv run python src/train.py ``` + +Or use a custom config file: + +```bash +uv run python src/train.py --config configs/my_experiment.yaml +``` + +Override specific parameters: + +```bash +uv run python src/train.py --learning-rate 0.001 --num-envs 32 --track +``` + +### Logging + +The training script uses a unified logging framework that: +- Logs to **WandB** (when enabled) +- Saves metrics to **local disk** (JSON files in `runs/`) +- Displays progress in **stdout** + +All experiment data is preserved locally, even if WandB is unavailable. + +## Project Structure + +``` +src/brittle_star_project/ # Core library (reusable components) +├── logging/ # Unified logging framework +├── environment/ # Environment wrappers +├── rl/ # RL algorithms and models +└── dataclasses/ # Configuration dataclasses + +configs/ # Training configurations +runs/ # Training outputs (checkpoints, metrics) +``` + +## For Researchers + +**Important:** Do not commit your personal WandB credentials to the repository. +Instead, create your own config file (e.g., `configs/yourname.yaml`) and add it to `.gitignore` if needed. + +See [configs/README.md](configs/README.md) for more details on configuration management. diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md index efd1399..ee9b4dc 100644 --- a/docs/CONTRIBUTING.md +++ b/docs/CONTRIBUTING.md @@ -33,3 +33,12 @@ Code readability is paramount, as code is read far more frequently than it is wr * **Simulation:** The simulation environment utilizes a MuJoCo brittle star. XML MuJoCo structures must remain realistic and respect morphological constraints. * **Experiment Tracking:** Weights & Biases (wandb) must be utilized for tracking and logging all experiments. * **Code Styling:** All code must conform to the chosen style guide (i.e. Google standard). This is enforced using build tools and pre-commit hooks such as flake8, black, or isort. + +## 5. AI-Assisted Development & Code Review + +This project supports AI-assisted development to enhance productivity, but contributors must take full responsibility for all AI-generated outputs. + +* **Self-Review Requirement:** Contributors must thoroughly self-review all AI-assisted code, documentation, and configurations before requesting peer review. This includes verifying correctness, adherence to project standards, scientific validity, and integration with existing code. +* **Quality Standards:** AI-generated content must meet the same rigorous standards as manually written code, including proper testing, documentation, and alignment with the scientific methodology outlined in Section 1. +* **Available Skills:** This project provides specific AI skills for common tasks (located in `.agents/skills/`), including linting and testing workflows. Contributors should leverage these skills to maintain consistency and quality. +* **Transparency:** When using AI assistance for complex algorithmic decisions or scientific design choices, contributors should document the rationale in commit messages or code comments where appropriate. From dc3531071fdb3acdaeb7e16ae37c5f26aead25a5 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 20:04:13 +0000 Subject: [PATCH 08/31] fix(devcontainer): update GPU runtime args for compatibility Change from deprecated --gpus flag to --device for better compatibility with newer Docker/CDI configurations. --- .devcontainer/devcontainer.json | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index be1fa12..3e5c8d4 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -34,8 +34,7 @@ }, "remoteUser": "vscode", "runArgs": [ - "--gpus", - "all" + "--device", "nvidia.com/gpu=all" ], // Ensure the .venv persists using a named volume for performance and parity "mounts": [ From 8ec693f04c35aadd12a11a9f372b36814b690ed1 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 20:52:51 +0000 Subject: [PATCH 09/31] fix(logging): improve JSON serialization and add wandb directory to gitignore - Fix float32 serialization issue in unified logger metrics flushing - Add jax.numpy import for proper type handling - Add wandb/ directory to .gitignore to exclude temporary tracking files - Tested wandb integration: metrics, artifacts, and local backup working correctly --- .gitignore | 3 +++ src/experiment_logger/unified_logger.py | 12 +++++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 8204ce9..6cc1e08 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,9 @@ artifacts/* runs/* +# Experiment tracking +wandb/ + # Python-generated files __pycache__/ *.py[oc] diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 87e3006..1501bf5 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -13,6 +13,7 @@ from pathlib import Path from typing import Any, Dict, Optional import flax +import jax.numpy as jnp import numpy as np from experiment_logger.wandb_utils import finish_wandb, init_wandb @@ -153,7 +154,16 @@ class UnifiedLogger: metrics_file = self.metrics_dir / "metrics.jsonl" with open(metrics_file, "a") as f: for metric in self.metrics_buffer: - f.write(json.dumps(metric) + "\n") + # Convert numpy/jax types to native Python types for JSON serialization + serializable_metric = {} + for k, v in metric.items(): + if hasattr(v, "item"): # numpy/jax scalar + serializable_metric[k] = v.item() + elif isinstance(v, (np.ndarray, jnp.ndarray)): + serializable_metric[k] = v.tolist() + else: + serializable_metric[k] = v + f.write(json.dumps(serializable_metric) + "\n") self.metrics_buffer.clear() except Exception as e: logger.error(f"Error flushing metrics: {e}") From d27a617199bbc08a790870f94cea7f03a33c2518 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 23:05:23 +0200 Subject: [PATCH 10/31] feat(experiment-logger): add config_utils module - Add config_utils.py: load_yaml_config, save_yaml_config, dataclass_from_dict, merge_config_with_cli, print_config - Export new symbols from package __init__.py --- src/experiment_logger/__init__.py | 3 +- src/experiment_logger/config_utils.py | 146 ++++++++++++++++++++++++++ 2 files changed, 148 insertions(+), 1 deletion(-) create mode 100644 src/experiment_logger/config_utils.py diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index 7ac168d..afbb13b 100644 --- a/src/experiment_logger/__init__.py +++ b/src/experiment_logger/__init__.py @@ -4,8 +4,9 @@ This package provides a unified interface for logging to multiple backends (WandB, disk, stdout) simultaneously, ensuring no data loss. """ +from experiment_logger.config_utils import load_yaml_config, merge_config_with_cli from experiment_logger.unified_logger import UnifiedLogger from experiment_logger.wandb_utils import finish_wandb, init_wandb -__all__ = ["UnifiedLogger", "init_wandb", "finish_wandb"] +__all__ = ["UnifiedLogger", "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 new file mode 100644 index 0000000..b296266 --- /dev/null +++ b/src/experiment_logger/config_utils.py @@ -0,0 +1,146 @@ +"""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__) + +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: + config = yaml.safe_load(f) + + if config is None: + return {} + + log.info(f"Loaded configuration from: {config_path}") + return config + + +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: + yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False) + + log.info(f"Saved configuration to: {config_path}") + + +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(): + if key in field_map: + field = field_map[key] + # Handle type conversion if needed + try: + # 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): + # 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') + else: + filtered_config[key] = field.type(value) if value is not None else None + except (ValueError, TypeError) as e: + log.warning(f"Could not convert {key}={value} to {field.type}: {e}") + 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: + config_idx = sys.argv.index("--config") + if config_idx + 1 < len(sys.argv): + extracted_config_file = sys.argv[config_idx + 1] + # 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): + yaml_config = load_yaml_config(extracted_config_file) + 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] + 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 + if yaml_value != default_value and yaml_value != cli_value: + log.info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})") + else: + 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) + + +def print_config(config: Any, title: str = "Configuration"): + """Pretty print configuration.""" + log.info(f"{title}:") + if is_dataclass(config): + for field in fields(config): + value = getattr(config, field.name) + 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 From 2001a92e75aa2d11764988f8584fb929e4e70def Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 23:05:33 +0200 Subject: [PATCH 11/31] feat(config): add YAML config templates for dev, production and personal use --- configs/README.md | 96 +++++++++++++++++++++++--------- configs/dev_test.yaml | 42 ++++++++++++++ configs/personal_template.yaml | 40 +++++++++++++ configs/production_training.yaml | 43 ++++++++++++++ 4 files changed, 195 insertions(+), 26 deletions(-) create mode 100644 configs/dev_test.yaml create mode 100644 configs/personal_template.yaml create mode 100644 configs/production_training.yaml diff --git a/configs/README.md b/configs/README.md index d06fbb2..ac6d5e1 100644 --- a/configs/README.md +++ b/configs/README.md @@ -2,47 +2,91 @@ This directory contains configuration files for training experiments. -## Usage +## Quick Start -Configuration files use YAML format and allow you to specify all training parameters in one place. +### 1. Choose a Template -### Quick Start - -Copy the default configuration template: +**For Development/Testing:** ```bash -cp configs/default_ppo.yaml configs/my_experiment.yaml +cp configs/dev_test.yaml configs/my_dev.yaml ``` -Edit `my_experiment.yaml` to customize your experiment settings, particularly: -- `wandb_entity`: Your WandB username or team name -- `track`: Set to `true` to enable WandB logging -- Training hyperparameters as needed +**For Production Training:** +```bash +cp configs/production_training.yaml configs/my_experiment.yaml +``` -Run training with your config: +### 2. Configure Your Settings + +Edit your config file and **set your wandb entity**: +```yaml +# ⚠️ IMPORTANT: Set this to your WandB username or team name +wandb_entity: "your-wandb-username" +track: true # Enable WandB logging +``` + +### 3. Run Training + +**Using config file:** ```bash python src/train.py --config configs/my_experiment.yaml ``` -### Override Parameters - -You can override any parameter from the command line: +**Override specific parameters:** ```bash python src/train.py --config configs/my_experiment.yaml --learning-rate 0.001 --num-envs 32 ``` -### Configuration for Different Users - -Each researcher should create their own config file with their WandB settings: -```yaml -# configs/researcher_name.yaml -exp_name: "researcher_name_experiment" -track: true -wandb_project_name: "PPO-Modularity" -wandb_entity: "your-wandb-username" # Change this! +**Pure CLI (no config file):** +```bash +python src/train.py --track --wandb-entity your-username --total-timesteps 1000000 ``` -This approach allows everyone to use the codebase without modifying source files. +## Features -## Available Configurations +### 📊 WandB Integration +- Real-time metrics logging +- Model checkpoints as artifacts +- Run comparison and collaboration -- `default_ppo.yaml` - Default PPO training configuration template +### 🔧 Flexible Configuration +- YAML files for reproducible experiments +- CLI overrides for quick adjustments +- Team collaboration without code changes + +## Configuration Templates + +### `dev_test.yaml` +- Fast iteration for development +- Short runs (100K timesteps) +- Frequent checkpoints +- Small environment count + +### `production_training.yaml` +- Full-scale training (50M timesteps) +- Optimized hyperparameters +- Production-ready settings + +### `default_ppo.yaml` +- Baseline configuration template +- Balanced settings for most use cases + +## Team Collaboration + +Each team member should create their own config file: + +```yaml +# configs/alice_experiment.yaml +exp_name: "alice_locomotion_v2" +track: true +wandb_project_name: "PPO-Modularity" +wandb_entity: "alice-research" # Alice's WandB username +total_timesteps: 20000000 +# ... other settings +``` + +This allows everyone to: +- Use their own WandB account +- Run different experiments simultaneously +- Share configurations via version control +- Avoid conflicts in run names diff --git a/configs/dev_test.yaml b/configs/dev_test.yaml new file mode 100644 index 0000000..194c544 --- /dev/null +++ b/configs/dev_test.yaml @@ -0,0 +1,42 @@ +# Quick Development/Testing Configuration +# +# Fast configuration for development and testing with short runs. + +# Experiment settings +exp_name: "brittle_star_dev_test" +seed: 123 + +# Tracking settings - IMPORTANT: Set your own wandb_entity! +track: true +wandb_project_name: "PPO-Modularity-Dev" +wandb_entity: null # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM + +# Model saving +save_model: true +checkpoint_frequency: 10 # More frequent checkpoints for testing + +# Environment settings +num_envs: 4 # Smaller for faster iteration + +# Training hyperparameters - Fast/testing +total_timesteps: 100000 # Short run for testing +learning_rate: 0.001 # Higher learning rate for faster learning +num_steps: 64 # Shorter rollouts +anneal_lr: true + +# PPO specific - Optimized for quick results +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 2 +update_epochs: 2 # Fewer epochs for speed +norm_adv: true +clip_coef: 0.1 +clip_vloss: true +ent_coef: 0.02 # Higher entropy for exploration +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null + +# Hardware +cuda: true +torch_deterministic: true \ No newline at end of file diff --git a/configs/personal_template.yaml b/configs/personal_template.yaml new file mode 100644 index 0000000..91c1ce6 --- /dev/null +++ b/configs/personal_template.yaml @@ -0,0 +1,40 @@ +# Personal Configuration Example for Team Member +# +# Copy this template and customize for your personal experiments + +# Experiment settings - PERSONALIZE THESE +exp_name: "YOUR_NAME_experiment_v1" # ⚠️ Change YOUR_NAME +seed: 42 + +# WandB settings - ⚠️ IMPORTANT: Set your credentials! +track: true # Enable WandB tracking +wandb_project_name: "PPO-Modularity" +wandb_entity: "YOUR_WANDB_USERNAME" # ⚠️ CHANGE THIS to your WandB username/team + +# Quick experiment settings (modify as needed) +total_timesteps: 500000 # 500K for quick results +num_envs: 8 +learning_rate: 0.0005 +num_steps: 128 + +# Model saving +save_model: true +checkpoint_frequency: 25 # Save checkpoints frequently + +# Standard PPO settings (usually don't need to change) +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 4 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null +anneal_lr: true + +# Hardware +cuda: true +torch_deterministic: true \ No newline at end of file diff --git a/configs/production_training.yaml b/configs/production_training.yaml new file mode 100644 index 0000000..54cf4ef --- /dev/null +++ b/configs/production_training.yaml @@ -0,0 +1,43 @@ +# Production Training Configuration +# +# Full-scale training configuration for production runs +# with wandb logging enabled. + +# Experiment settings +exp_name: "brittle_star_production" +seed: 42 + +# Tracking settings - IMPORTANT: Set your own wandb_entity! +track: true +wandb_project_name: "PPO-Modularity" +wandb_entity: null # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM + +# Model saving +save_model: true +checkpoint_frequency: 100 # Save checkpoint every 100 iterations + +# Environment settings +num_envs: 32 # Increased for production + +# Training hyperparameters - Production scale +total_timesteps: 50000000 # 50M timesteps for full training +learning_rate: 0.00025 +num_steps: 256 # Longer rollouts +anneal_lr: true + +# PPO specific - Fine-tuned +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 8 # More minibatches for stability +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null + +# Hardware +cuda: true +torch_deterministic: true \ No newline at end of file From e160a55d95c6ec829fc6acf4c9428b417ecd0477 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 23:05:43 +0200 Subject: [PATCH 12/31] fix(train): integrate YAML config --- pyproject.toml | 1 + src/brittle_star_project/dataclasses/PPOArgs.py | 2 +- src/train.py | 9 ++++++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 43c977b..9ba5a9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "optax>=0.2.6", "pyopengl>=3.1.10", "pyopengl-accelerate>=3.1.10", + "pyyaml>=6.0", "tyro>=1.0.10", "wandb==0.24.2", ] diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index b27265f..3a6836b 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/src/brittle_star_project/dataclasses/PPOArgs.py @@ -26,7 +26,7 @@ class PPOArgs: wandb_project_name: str = "PPO-Modularity" # the entity (team) of wandb's project - wandb_entity: str | None = "tdpeuter-ghent-university" + wandb_entity: str | None = None # whether to capture videos of the agent performances (check out `videos` folder) capture_video: bool = False diff --git a/src/train.py b/src/train.py index 00e8c60..ad31b56 100644 --- a/src/train.py +++ b/src/train.py @@ -21,6 +21,7 @@ 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__) @@ -418,7 +419,13 @@ def main() -> None: level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) - args = tyro.cli(PPOArgs) + + # Enhanced argument parsing with YAML config support + args = merge_config_with_cli(PPOArgs) + + # Print final configuration + print_config(args, "Final Training Configuration") + train(args) From 8347d81d705f40297822250fc101c649216ab17c Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 23:05:52 +0200 Subject: [PATCH 13/31] test: add tests for YAML config loading - test_config.py: covers load_yaml_config (happy path, missing file raises) - Uses project-relative paths so tests run in any environment --- tests/test_config.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tests/test_config.py diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..212e4b4 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,38 @@ +"""Tests for YAML config loading.""" + +import sys +from pathlib import Path + +import pytest + +# Ensure src is on the path when running from the project root +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +CONFIGS_DIR = Path(__file__).parent.parent / "configs" + + +class TestYamlConfig: + def test_load_yaml_config(self): + from experiment_logger.config_utils import load_yaml_config + + config = load_yaml_config(str(CONFIGS_DIR / "default_ppo.yaml")) + assert isinstance(config, dict) + assert "total_timesteps" in config + assert "learning_rate" in config + + def test_load_dev_test_config(self): + from experiment_logger.config_utils import load_yaml_config + + config = load_yaml_config(str(CONFIGS_DIR / "dev_test.yaml")) + assert config["total_timesteps"] == 100000 + + def test_missing_config_raises(self): + from experiment_logger.config_utils import load_yaml_config + + with pytest.raises(FileNotFoundError): + load_yaml_config("nonexistent.yaml") + + def test_merge_config_with_cli_is_callable(self): + from experiment_logger.config_utils import merge_config_with_cli + + assert callable(merge_config_with_cli) From 3e992a33774f1eb723fc4cf77cb4db2fb141b6c8 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 1 Apr 2026 07:59:26 +0000 Subject: [PATCH 14/31] chore: update uv.lock --- uv.lock | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/uv.lock b/uv.lock index e59c767..679acc3 100644 --- a/uv.lock +++ b/uv.lock @@ -25,6 +25,7 @@ dependencies = [ { name = "optax" }, { name = "pyopengl" }, { name = "pyopengl-accelerate" }, + { name = "pyyaml" }, { name = "tyro" }, { name = "wandb" }, ] @@ -56,6 +57,7 @@ requires-dist = [ { name = "optax", specifier = ">=0.2.6" }, { name = "pyopengl", specifier = ">=3.1.10" }, { name = "pyopengl-accelerate", specifier = ">=3.1.10" }, + { name = "pyyaml", specifier = ">=6.0" }, { name = "tyro", specifier = ">=1.0.10" }, { name = "wandb", specifier = "==0.24.2" }, ] @@ -1630,7 +1632,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ From 7e7c5bf27cec1e55db4ae4101769283a1d7ad115 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 1 Apr 2026 10:26:39 +0200 Subject: [PATCH 15/31] refactor: set wandb entity and simplify READMEs --- configs/README.md | 89 ++--------- configs/default_ppo.yaml | 2 +- configs/dev_test.yaml | 2 +- configs/personal_template.yaml | 2 +- configs/production_training.yaml | 2 +- .../dataclasses/PPOArgs.py | 2 +- src/experiment_logger/README.md | 140 +----------------- 7 files changed, 21 insertions(+), 218 deletions(-) diff --git a/configs/README.md b/configs/README.md index ac6d5e1..c1662fb 100644 --- a/configs/README.md +++ b/configs/README.md @@ -2,91 +2,22 @@ This directory contains configuration files for training experiments. -## Quick Start +## Usage -### 1. Choose a Template +Use `--config` with `src/train.py` to run an experiment: -**For Development/Testing:** ```bash -cp configs/dev_test.yaml configs/my_dev.yaml +python src/train.py --config configs/default_ppo.yaml ``` -**For Production Training:** +You can overriding settings via CLI: ```bash -cp configs/production_training.yaml configs/my_experiment.yaml +python src/train.py --config configs/default_ppo.yaml --learning-rate 0.001 ``` -### 2. Configure Your Settings +## Available Configurations -Edit your config file and **set your wandb entity**: -```yaml -# ⚠️ IMPORTANT: Set this to your WandB username or team name -wandb_entity: "your-wandb-username" -track: true # Enable WandB logging -``` - -### 3. Run Training - -**Using config file:** -```bash -python src/train.py --config configs/my_experiment.yaml -``` - -**Override specific parameters:** -```bash -python src/train.py --config configs/my_experiment.yaml --learning-rate 0.001 --num-envs 32 -``` - -**Pure CLI (no config file):** -```bash -python src/train.py --track --wandb-entity your-username --total-timesteps 1000000 -``` - -## Features - -### 📊 WandB Integration -- Real-time metrics logging -- Model checkpoints as artifacts -- Run comparison and collaboration - -### 🔧 Flexible Configuration -- YAML files for reproducible experiments -- CLI overrides for quick adjustments -- Team collaboration without code changes - -## Configuration Templates - -### `dev_test.yaml` -- Fast iteration for development -- Short runs (100K timesteps) -- Frequent checkpoints -- Small environment count - -### `production_training.yaml` -- Full-scale training (50M timesteps) -- Optimized hyperparameters -- Production-ready settings - -### `default_ppo.yaml` -- Baseline configuration template -- Balanced settings for most use cases - -## Team Collaboration - -Each team member should create their own config file: - -```yaml -# configs/alice_experiment.yaml -exp_name: "alice_locomotion_v2" -track: true -wandb_project_name: "PPO-Modularity" -wandb_entity: "alice-research" # Alice's WandB username -total_timesteps: 20000000 -# ... other settings -``` - -This allows everyone to: -- Use their own WandB account -- Run different experiments simultaneously -- Share configurations via version control -- Avoid conflicts in run names +- `default_ppo.yaml`: Baseline config. +- `dev_test.yaml`: Fast iteration for development. +- `production_training.yaml`: Full-scale training. +- `personal_template.yaml`: Template for team members to customize. diff --git a/configs/default_ppo.yaml b/configs/default_ppo.yaml index 5775a84..1b06c3d 100644 --- a/configs/default_ppo.yaml +++ b/configs/default_ppo.yaml @@ -15,7 +15,7 @@ seed: 1 # Tracking settings track: false # Set to true to enable WandB logging wandb_project_name: "PPO-Modularity" -wandb_entity: null # Set to your WandB username or team name +wandb_entity: "SEL3-2026-Groep-4" # Set to your WandB username or team name # Model saving save_model: true diff --git a/configs/dev_test.yaml b/configs/dev_test.yaml index 194c544..b64d330 100644 --- a/configs/dev_test.yaml +++ b/configs/dev_test.yaml @@ -9,7 +9,7 @@ seed: 123 # Tracking settings - IMPORTANT: Set your own wandb_entity! track: true wandb_project_name: "PPO-Modularity-Dev" -wandb_entity: null # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM +wandb_entity: "SEL3-2026-Groep-4" # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM # Model saving save_model: true diff --git a/configs/personal_template.yaml b/configs/personal_template.yaml index 91c1ce6..67caef7 100644 --- a/configs/personal_template.yaml +++ b/configs/personal_template.yaml @@ -9,7 +9,7 @@ seed: 42 # WandB settings - ⚠️ IMPORTANT: Set your credentials! track: true # Enable WandB tracking wandb_project_name: "PPO-Modularity" -wandb_entity: "YOUR_WANDB_USERNAME" # ⚠️ CHANGE THIS to your WandB username/team +wandb_entity: "SEL3-2026-Groep-4" # ⚠️ CHANGE THIS to your WandB username/team # Quick experiment settings (modify as needed) total_timesteps: 500000 # 500K for quick results diff --git a/configs/production_training.yaml b/configs/production_training.yaml index 54cf4ef..351ee83 100644 --- a/configs/production_training.yaml +++ b/configs/production_training.yaml @@ -10,7 +10,7 @@ seed: 42 # Tracking settings - IMPORTANT: Set your own wandb_entity! track: true wandb_project_name: "PPO-Modularity" -wandb_entity: null # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM +wandb_entity: "SEL3-2026-Groep-4" # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM # Model saving save_model: true diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 3a6836b..d59defa 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/src/brittle_star_project/dataclasses/PPOArgs.py @@ -26,7 +26,7 @@ class PPOArgs: wandb_project_name: str = "PPO-Modularity" # the entity (team) of wandb's project - wandb_entity: str | None = None + wandb_entity: str | None = "SEL3-2026-Groep-4" # whether to capture videos of the agent performances (check out `videos` folder) capture_video: bool = False diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md index 2beac1a..ab6394d 100644 --- a/src/experiment_logger/README.md +++ b/src/experiment_logger/README.md @@ -1,145 +1,17 @@ # Experiment Logger -A lightweight, standalone logging framework for machine learning experiments with multi-backend support. +A lightweight logging framework supporting Weights & Biases, local JSON, and stdout. -## Features - -- **Multi-backend logging**: Simultaneously log to WandB, local disk (JSON), and stdout -- **Data preservation**: All metrics saved locally, even if WandB is unavailable -- **Checkpoint management**: Save model checkpoints with metadata -- **WandB integration**: Optional artifact upload for model versioning -- **Graceful degradation**: Works without WandB installed -- **Simple API**: Minimal configuration required - -## Installation - -This package is included in the project. To use it in your code: - -```python -from experiment_logger import UnifiedLogger -``` - -## Quick Start +## Usage ```python from experiment_logger import UnifiedLogger -# Initialize logger -logger = UnifiedLogger( - run_name="my_experiment", - config={"learning_rate": 0.001, "batch_size": 32}, - project_name="MyProject", - entity="my-wandb-username", # Optional - use_wandb=True, # Set to False to disable WandB -) +logger = UnifiedLogger(run_name="my_experiment", config={"lr": 0.001}) -# Log metrics -for step in range(100): - logger.log({ - "loss": 1.0 / (step + 1), - "accuracy": step * 0.01, - }, step=step) - -# Save checkpoint -logger.save_checkpoint( - params=model_params, - step=100, - metadata={"epoch": 1, "val_acc": 0.95}, -) - -# Save final model -logger.save_final_model( - params=final_params, - metadata={"final_accuracy": 0.98}, -) - -# Finalize (flushes remaining metrics) +logger.log({"loss": 0.5}, step=1) +logger.save_checkpoint(params=model_params, step=1) logger.finish() ``` -## Context Manager - -Use as a context manager for automatic cleanup: - -```python -with UnifiedLogger(run_name="my_exp", config={}) as logger: - logger.log({"metric": 1.0}) - # Automatically calls finish() on exit -``` - -## Configuration - -### Constructor Parameters - -- `run_name` (str): Unique name for this run -- `config` (dict): Configuration dictionary with hyperparameters -- `project_name` (str): WandB project name (default: "PPO-Modularity") -- `entity` (str, optional): WandB entity (team/user name) -- `base_dir` (str): Base directory for local storage (default: "runs") -- `use_wandb` (bool): Enable WandB logging (default: True) -- `save_code` (bool): Save code to WandB (default: True) - -### Directory Structure - -``` -runs/ -└── my_experiment/ - ├── config.json # Saved configuration - ├── metrics/ - │ └── metrics.jsonl # Line-delimited JSON metrics - ├── checkpoints/ - │ ├── checkpoint_step_100.flax - │ └── checkpoint_step_100_metadata.json - └── final_model.flax -``` - -## API Reference - -### `log(metrics, step=None, commit=True)` - -Log metrics to all backends. - -**Parameters:** -- `metrics` (dict): Dictionary of metric name -> value -- `step` (int, optional): Global step counter (auto-incremented if None) -- `commit` (bool): Whether to commit to WandB immediately - -### `save_checkpoint(params, step, prefix="checkpoint", metadata=None)` - -Save model checkpoint to disk and optionally to WandB. - -**Parameters:** -- `params`: Model parameters (Flax params or any serializable object) -- `step` (int): Current training step -- `prefix` (str): Prefix for checkpoint filename -- `metadata` (dict, optional): Additional metadata to save - -### `save_final_model(params, metadata=None)` - -Save the final trained model. - -**Parameters:** -- `params`: Model parameters -- `metadata` (dict, optional): Metadata about the final model - -### `finish()` - -Finalize logging and cleanup. Flushes remaining metrics to disk. - -## Usage in Projects - -This logger is designed to be: -- **Project-agnostic**: Use in any ML project, not just this one -- **Framework-agnostic**: Works with JAX, PyTorch, TensorFlow, etc. -- **Minimal dependencies**: Only requires `wandb` (optional), `flax` (for serialization), and `numpy` - -## Design Philosophy - -1. **Never lose data**: All metrics saved locally, regardless of WandB availability -2. **Simple API**: Minimal boilerplate, easy to integrate -3. **Fail gracefully**: Missing WandB shouldn't break experiments -4. **Reproducibility**: Save full configuration with every run - -## License - -Part of the 2026SEL3-project-BrittleStar repository. +Logs and checkoints are saved in the `runs/` directory. If `track=True` (or `use_wandb=True`), everything is additionally synced to Weights & Biases. From ff901013774430569ec75896af36623af96c5d23 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 1 Apr 2026 16:01:12 +0000 Subject: [PATCH 16/31] refactor(log): improved logging workflow --- .../environment/BrittleStarJaxEnvWrapper.py | 8 ++ .../environment/factory.py | 8 +- src/experiment_logger/__init__.py | 11 +- src/experiment_logger/config_utils.py | 57 ++++---- src/experiment_logger/unified_logger.py | 128 +++++++++++++----- src/main.py | 7 +- src/train.py | 33 ++--- 7 files changed, 162 insertions(+), 90 deletions(-) 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) From d4ac45b34e7bae50a72e25251e67773311860372 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 18:53:31 +0200 Subject: [PATCH 17/31] feat(logger): updated run_name format and added runtime log level filtering --- src/experiment_logger/unified_logger.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index c391f52..15c8a19 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -6,6 +6,7 @@ This logger ensures all experimental data is preserved by writing to: 3. stdout (for real-time monitoring) """ +import datetime import json import logging import subprocess @@ -38,8 +39,8 @@ def get_logger() -> "UnifiedLogger": except Exception: commit_hash = "unknown" - timestamp = int(time.time()) - generic_name = f"brittle_star_{commit_hash}_{timestamp}" + timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + generic_name = f"{timestamp_str}_{commit_hash}_brittle_star" # Initialize generic fallback logger without WandB _global_logger = UnifiedLogger( @@ -65,6 +66,7 @@ class UnifiedLogger: base_dir: str = "runs", use_wandb: bool = True, save_code: bool = True, + log_level: int = logging.INFO, _set_as_global: bool = True, ): """Initialize the unified logger. @@ -100,7 +102,7 @@ class UnifiedLogger: # 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) + self._text_logger.setLevel(log_level) # Avoid duplicate handlers if re-instantiated if not self._text_logger.handlers: @@ -134,6 +136,10 @@ class UnifiedLogger: self.info(f"Local storage: {self.run_dir.absolute()}") self.info(f"WandB logging: {self.wandb_available}") + def set_level(self, level: int): + """Dynamically update the verbosity of the stdout/text logger.""" + self._text_logger.setLevel(level) + def info(self, msg: str, *args, **kwargs): """Log an info message to stdout and disk.""" self._text_logger.info(msg, *args, **kwargs) From 795ae48520740679709c0027407cc7c1d473ce25 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 18:58:33 +0200 Subject: [PATCH 18/31] feat(log): add TensorBoard support and enhance WandB API key checking --- src/experiment_logger/unified_logger.py | 23 +++++++++++++++++++++++ src/experiment_logger/wandb_utils.py | 21 +++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 15c8a19..06ca3a9 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -124,6 +124,16 @@ class UnifiedLogger: # Save config to disk self._save_config() + # Setup TensorBoard + self.writer = None + try: + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(self.run_dir) + self.info("TensorBoard SummaryWriter initialized.") + except ImportError: + self.warning("tensorboard not installed. Skipping SummaryWriter.") + # Initialize WandB if requested if self.use_wandb: self._init_wandb(project_name, entity, save_code) @@ -206,6 +216,16 @@ class UnifiedLogger: except Exception as e: self.warning(f"WandB logging failed: {e}") + # Log to TensorBoard + if self.writer is not None: + for k, v in metrics.items(): + if isinstance(v, (int, float, np.floating, np.integer)): + self.writer.add_scalar(k, v, step) + elif hasattr(v, "item"): + self.writer.add_scalar(k, v.item(), step) + elif isinstance(v, (np.ndarray, jnp.ndarray)) and v.size == 1: + self.writer.add_scalar(k, v.item(), step) + # Buffer for disk storage self.metrics_buffer.append(metrics_with_metadata) @@ -331,6 +351,9 @@ class UnifiedLogger: # Flush remaining metrics self._flush_metrics() + if self.writer is not None: + self.writer.close() + self.info(f"Run complete. Results saved to: {self.run_dir.absolute()}") # Finish WandB run diff --git a/src/experiment_logger/wandb_utils.py b/src/experiment_logger/wandb_utils.py index 5fd8837..2c162fb 100644 --- a/src/experiment_logger/wandb_utils.py +++ b/src/experiment_logger/wandb_utils.py @@ -36,6 +36,27 @@ def init_wandb( """ try: import wandb + import os + import sys + + # Robust HPC checking: check for API key + has_key = os.environ.get("WANDB_API_KEY") is not None + if not has_key: + try: + # Check if logged in locally via settings/netrc + has_key = wandb.setup().settings.api_key is not None + except Exception: + pass + + is_interactive = sys.stdout.isatty() + + if not has_key and not is_interactive and os.environ.get("WANDB_MODE") != "offline": + logger.warning( + "WANDB_API_KEY not found and environment is non-interactive. Switching to offline mode." + ) + sync_path = f"runs/{name}" if name else "runs" + logger.warning(f"WandB is offline. Use 'wandb sync {sync_path}' to upload logs later.") + os.environ["WANDB_MODE"] = "offline" run = wandb.init( project=project, From faf31567e8bb526d673fe7fbcc0ad2efdcf5ef51 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 19:32:48 +0200 Subject: [PATCH 19/31] refactor(log): PPOTrainer --- .../trainers/PPOTrainer.py | 168 +++++++----------- 1 file changed, 61 insertions(+), 107 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 7df57f0..0108e7f 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -13,7 +13,8 @@ import numpy as np import optax import tqdm from flax.training.train_state import TrainState -from torch.utils.tensorboard import SummaryWriter + +from experiment_logger import get_logger from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper @@ -207,7 +208,7 @@ class PPOTrainer: self.env = env self.run_dir = run_dir self.run_name = run_name - self.writer = SummaryWriter(self.run_dir) + self.logger = get_logger() self.key = jax.random.PRNGKey(args.seed) @@ -247,16 +248,14 @@ class PPOTrainer: self._init_random() - def _init_random(self, log: bool = True): - if log: - print(f"[RANDOM]: Setting random seed to {self.args.seed}") + def _init_random(self): + self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}") random.seed(self.args.seed) np.random.seed(self.args.seed) - def _init_agent(self, log: bool = True): - if log: - print("[AGENT]: Initializing agent...") + def _init_agent(self): + self.logger.info("[AGENT]: Initializing agent...") sensor = GenericDenseLayersWithActivation() feature_extractor = GenericDenseLayersWithActivation() @@ -267,9 +266,8 @@ class PPOTrainer: # messenger = OneDenseLayerMLP() return sensor, feature_extractor, actor, critic - def _init_agent_state(self, log: bool = True) -> TrainState: - if log: - print("[AGENT STATE]: Initializing agent state...") + def _init_agent_state(self) -> TrainState: + self.logger.info("[AGENT STATE]: Initializing agent state...") self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split( self.key, 5 @@ -313,9 +311,8 @@ class PPOTrainer: ), ) - def _init_episode_stats(self, log: bool = True) -> EpisodeStatistics: - if log: - print("[EPISODE STATS]: Initializing episode stats...") + def _init_episode_stats(self) -> EpisodeStatistics: + self.logger.info("[EPISODE STATS]: Initializing episode stats...") return EpisodeStatistics( episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32), @@ -350,39 +347,29 @@ class PPOTrainer: iteration_time_start, loss_info, ): + metrics = { + "charts/avg_episodic_return": loss_info.avg_episodic_return, + "charts/avg_episodic_length": np.mean( + jax.device_get(episode_stats.returned_episode_lengths) + ), + "charts/learning_rate": self.agent_state.opt_state[1] + .hyperparams["learning_rate"] + .item(), + "losses/value_loss": loss_info.v_loss[-1, -1].item(), + "losses/policy_loss": loss_info.pg_loss[-1, -1].item(), + "losses/entropy": loss_info.entropy_loss[-1, -1].item(), + "losses/approx_kl": loss_info.approx_kl[-1, -1].item(), + "losses/loss": loss_info.loss[-1, -1].item(), + "charts/SPS": int(global_step / (time.time() - start_time)), + "charts/SPS_update": int( + self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start) + ), + } + self.logger.log(metrics, step=global_step) - self.writer.add_scalar( - "charts/avg_episodic_return", loss_info.avg_episodic_return, global_step - ) - self.writer.add_scalar( - "charts/avg_episodic_length", - np.mean(jax.device_get(episode_stats.returned_episode_lengths)), - global_step, - ) - self.writer.add_scalar( - "charts/learning_rate", - self.agent_state.opt_state[1].hyperparams["learning_rate"].item(), - global_step, - ) - self.writer.add_scalar("losses/value_loss", loss_info.v_loss[-1, -1].item(), global_step) - self.writer.add_scalar("losses/policy_loss", loss_info.pg_loss[-1, -1].item(), global_step) - self.writer.add_scalar("losses/entropy", loss_info.entropy_loss[-1, -1].item(), global_step) - self.writer.add_scalar("losses/approx_kl", loss_info.approx_kl[-1, -1].item(), global_step) - self.writer.add_scalar("losses/loss", loss_info.loss[-1, -1].item(), global_step) - self.writer.add_scalar( - "charts/SPS", int(global_step / (time.time() - start_time)), global_step - ) - self.writer.add_scalar( - "charts/SPS_update", - int(self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)), - global_step, - ) - - def _step( - self, env_state, next_obs, next_done, is_tty: bool, iteration: int, log: bool = True - ) -> tuple: - if log and not is_tty and iteration == 1: - print(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}", flush=True) + def _step(self, env_state, next_obs, next_done, is_tty: bool, iteration: int) -> tuple: + if not is_tty and iteration == 1: + self.logger.info(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}") ( self.agent_state, @@ -394,20 +381,20 @@ class PPOTrainer: next_env_state, ) = self._rollout(env_state, next_obs, next_done) - if log and not is_tty and iteration == 1: - print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True) + if not is_tty and iteration == 1: + self.logger.info(f">>> [HPC] First rollout completed: {time.ctime()}") storage = self._compute_gae(storage, next_obs, next_done) - if log and not is_tty and iteration == 1: - print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True) + if not is_tty and iteration == 1: + self.logger.info(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}") self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = ( self._ppo.update_ppo(self.agent_state, storage, self.key) ) - if log and not is_tty and iteration == 1: - print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True) + if not is_tty and iteration == 1: + self.logger.info(f">>> [HPC] First PPO update completed: {time.ctime()}") avg_episodic_return = float( jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item() @@ -429,77 +416,45 @@ class PPOTrainer: def _close(self): self.env.close() - self.writer.close() - def _save_model(self, model_path: str, log: bool = True): - if log: - print(f"[SAVE]: Saving the model to: {model_path}...") + def _save_model(self, model_path: str): + self.logger.info("[SAVE]: Saving the final model...") - with open(model_path, "wb") as f: - f.write( - flax.serialization.to_bytes( - [ - vars(self.args), - [ - self.agent_state.params["sensor_params"], - self.agent_state.params["actor_params"], - self.agent_state.params["critic_params"], - self.agent_state.params["feature_extractor_params"], - ], - ] - ) - ) + params = [ + vars(self.args), + [ + self.agent_state.params["sensor_params"], + self.agent_state.params["actor_params"], + self.agent_state.params["critic_params"], + self.agent_state.params["feature_extractor_params"], + ], + ] + self.logger.save_final_model(params=params) - def train(self, log: bool = True): + def train(self): """ Train the PPO agent for a specified number of iterations (passed through PPOArgs in constructor). Closes the environment at the end of training. """ - if log: - print(f"running name: {self.run_name}") + self.logger.info(f"running name: {self.run_name}") is_tty = sys.stdout.isatty() - if log: - print("[TRAIN]: Resetting environment...") + self.logger.info("[TRAIN]: Resetting environment...") - if not is_tty: - print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True) + if not is_tty: + self.logger.info(f">>> [HPC] Initial reset started: {time.ctime()}") env_state = self.env.reset(seed=self.args.seed) next_obs = _convert_obs_dict_to_array(env_state.observations) next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_) - if log and not is_tty: - print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True) + if not is_tty: + self.logger.info(f">>> [HPC] Initial reset completed: {time.ctime()}") global_step = 0 start_time = time.time() - if self.args.track: - import wandb - - if log: - print("[TRAIN]: Initializing Weights and Biases...") - - wandb.init( - project=self.args.wandb_project_name, - entity=self.args.wandb_entity, - sync_tensorboard=True, - config=vars(self.args), - name=self.run_name, - save_code=True, - ) - - if log: - print("[TRAIN]: Adding hyperparameters to TensorBoard...") - - self.writer.add_text( - "hyperparameters", - "|param|value|\n|---|---|\n" - + "\n".join(f"|{k}|{v}|" for k, v in vars(self.args).items()), - ) - iter_bar = tqdm.tqdm( range(1, self.args.num_iterations + 1), disable=not is_tty, @@ -514,19 +469,18 @@ class PPOTrainer: global_step += self.args.num_steps * self.args.num_envs self._log(global_step, self.episode_stats, start_time, iteration_time_start, loss_info) - if log and not is_tty: + if not is_tty: sps = int(global_step / (time.time() - start_time)) remaining_steps = self.args.total_timesteps - global_step eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 eta_str = str(datetime.timedelta(seconds=eta_seconds)) - print( + self.logger.info( f"Iteration {iteration}/{self.args.num_iterations} | " f"Step {global_step}/{self.args.total_timesteps} | " f"SPS {sps} | " f"Return {loss_info.avg_episodic_return:.4f} | " - f"ETA {eta_str}", - flush=True, + f"ETA {eta_str}" ) if self.args.save_model: From 46ae13abc3f9fe772610514a384395ebaac962b2 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 19:54:05 +0200 Subject: [PATCH 20/31] feat(log): non-interactive logging, incl. progress bar --- .../trainers/PPOTrainer.py | 58 ++++++++----------- src/experiment_logger/unified_logger.py | 14 +++++ 2 files changed, 38 insertions(+), 34 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 0108e7f..1c807d4 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -1,6 +1,5 @@ import datetime import random -import sys import time from dataclasses import asdict, dataclass from functools import partial @@ -11,7 +10,6 @@ import jax import jax.numpy as jnp import numpy as np import optax -import tqdm from flax.training.train_state import TrainState from experiment_logger import get_logger @@ -367,9 +365,9 @@ class PPOTrainer: } self.logger.log(metrics, step=global_step) - def _step(self, env_state, next_obs, next_done, is_tty: bool, iteration: int) -> tuple: - if not is_tty and iteration == 1: - self.logger.info(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}") + def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple: + if iteration == 1: + self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}") ( self.agent_state, @@ -381,20 +379,20 @@ class PPOTrainer: next_env_state, ) = self._rollout(env_state, next_obs, next_done) - if not is_tty and iteration == 1: - self.logger.info(f">>> [HPC] First rollout completed: {time.ctime()}") + if iteration == 1: + self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}") storage = self._compute_gae(storage, next_obs, next_done) - if not is_tty and iteration == 1: - self.logger.info(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}") + if iteration == 1: + self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}") self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = ( self._ppo.update_ppo(self.agent_state, storage, self.key) ) - if not is_tty and iteration == 1: - self.logger.info(f">>> [HPC] First PPO update completed: {time.ctime()}") + if iteration == 1: + self.logger.log_non_interactive(f"First PPO update completed: {time.ctime()}") avg_episodic_return = float( jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item() @@ -439,49 +437,41 @@ class PPOTrainer: """ self.logger.info(f"running name: {self.run_name}") - is_tty = sys.stdout.isatty() self.logger.info("[TRAIN]: Resetting environment...") - - if not is_tty: - self.logger.info(f">>> [HPC] Initial reset started: {time.ctime()}") + self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}") env_state = self.env.reset(seed=self.args.seed) next_obs = _convert_obs_dict_to_array(env_state.observations) next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_) - if not is_tty: - self.logger.info(f">>> [HPC] Initial reset completed: {time.ctime()}") + self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}") global_step = 0 start_time = time.time() - iter_bar = tqdm.tqdm( - range(1, self.args.num_iterations + 1), - disable=not is_tty, - ) + iter_bar = self.logger.progress_bar(range(1, self.args.num_iterations + 1)) for iteration in iter_bar: iteration_time_start = time.time() env_state, next_obs, next_done, loss_info = self._step( - env_state, next_obs, next_done, is_tty=is_tty, iteration=iteration + env_state, next_obs, next_done, iteration=iteration ) global_step += self.args.num_steps * self.args.num_envs self._log(global_step, self.episode_stats, start_time, iteration_time_start, loss_info) - if not is_tty: - sps = int(global_step / (time.time() - start_time)) - remaining_steps = self.args.total_timesteps - global_step - eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 - eta_str = str(datetime.timedelta(seconds=eta_seconds)) + sps = int(global_step / (time.time() - start_time)) + remaining_steps = self.args.total_timesteps - global_step + eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 + eta_str = str(datetime.timedelta(seconds=eta_seconds)) - self.logger.info( - f"Iteration {iteration}/{self.args.num_iterations} | " - f"Step {global_step}/{self.args.total_timesteps} | " - f"SPS {sps} | " - f"Return {loss_info.avg_episodic_return:.4f} | " - f"ETA {eta_str}" - ) + self.logger.log_non_interactive( + f"Iteration {iteration}/{self.args.num_iterations} | " + f"Step {global_step}/{self.args.total_timesteps} | " + f"SPS {sps} | " + f"Return {loss_info.avg_episodic_return:.4f} | " + f"ETA {eta_str}" + ) if self.args.save_model: model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model" diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 06ca3a9..71380b0 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -10,6 +10,7 @@ import datetime import json import logging import subprocess +import sys import time from pathlib import Path from typing import Any, Dict, List, Optional @@ -86,6 +87,7 @@ class UnifiedLogger: self.use_wandb = use_wandb self.wandb_available = False self.wandb_run = None + self.is_interactive = sys.stdout.isatty() # Setup local storage self.run_dir = Path(base_dir) / run_name @@ -150,6 +152,18 @@ class UnifiedLogger: """Dynamically update the verbosity of the stdout/text logger.""" self._text_logger.setLevel(level) + def log_non_interactive(self, msg: str, *args, **kwargs): + """Log an info message only if running in a non-interactive environment.""" + if not self.is_interactive: + self.info(msg, *args, **kwargs) + + def progress_bar(self, iterable=None, *args, **kwargs): + """Wrapper around tqdm that automatically disables in non-interactive environments.""" + import tqdm + + kwargs.setdefault("disable", not self.is_interactive) + return tqdm.tqdm(iterable, *args, **kwargs) + def info(self, msg: str, *args, **kwargs): """Log an info message to stdout and disk.""" self._text_logger.info(msg, *args, **kwargs) From d9202b0390ac8f9c3cf0a948fdd9aa1d1a3e7d10 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 20:14:31 +0200 Subject: [PATCH 21/31] refactor(log): use UnifiedLogger in training script --- .env.example | 12 ++++++++++++ scripts/train.py | 42 ++++++++++++++++++++++-------------------- 2 files changed, 34 insertions(+), 20 deletions(-) create mode 100644 .env.example diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..ebd4eb9 --- /dev/null +++ b/.env.example @@ -0,0 +1,12 @@ +# Brittle Star Project Environment Variables +# Copy this file to .env and fill in your values. +# IMPORTANT: Never commit the actual .env file, it is in .gitignore + +# ---------------------------- # +# Weights and Biases API Key # +# ---------------------------- # +# To find your API key: +# 1. Log in to wandb.ai +# 2. Go to User Settings (https://wandb.ai/settings) +# 3. Scroll down to the "API keys" section +WANDB_API_KEY=your_api_key_here diff --git a/scripts/train.py b/scripts/train.py index 9ee6da6..39b062d 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -10,6 +10,9 @@ from brittle_star_project.dataclasses import PPOArgs from brittle_star_project.trainers.PPOTrainer import PPOTrainer from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from experiment_logger.unified_logger import get_logger +from experiment_logger.config_utils import merge_config_with_cli, print_config + def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper: if config_path is None: @@ -17,28 +20,15 @@ def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs) -def parse_args(log: bool = True) -> PPOArgs: - temp_args = tyro.cli(PPOArgs) +def parse_args() -> PPOArgs: + import argparse - if temp_args.hyperparameter_config_path is not None: - if log: - print(f"Loading hyperparameter config from {temp_args.hyperparameter_config_path}") + # Use argparse to reliably extract just the config path without swallowing --help + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--hyperparameter-config-path", type=str, default=None) + known_args, _ = parser.parse_known_args() - with open(temp_args.hyperparameter_config_path, "r") as f: - config = yaml.safe_load(f) - if config: - # parse PPOArgs with defaults from yaml. - for key, value in config.items(): - if hasattr(temp_args, key): - setattr(temp_args, key, value) - - # Reparse CLI to ensure they OVERRIDE the yaml - args = tyro.cli(PPOArgs, default=temp_args) - else: - if log: - print("No hyperparameter config provided, using default config") - - args = temp_args + args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path) return args @@ -60,6 +50,7 @@ if __name__ == "__main__": git_hash = get_git_hash() run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" + if args.run_dir is None: run_dir = f"runs/{run_name}" else: @@ -67,6 +58,17 @@ if __name__ == "__main__": os.makedirs(run_dir, exist_ok=True) + # Initialize Global Logger + logger = get_logger() + logger.init( + project_name=args.wandb_project_name, # or default PPO-Modularity if missing + run_name=run_name, + base_dir=os.path.dirname(run_dir), + use_wandb=args.track, + ) + + print_config(args, title="PPO Training Configuration") + env = make_env(args.env_config_path, args.num_envs) torch.backends.cudnn.deterministic = args.torch_deterministic From ae483b306fa97b0b06a4494992515ddb1ab31f8c Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 20:26:26 +0200 Subject: [PATCH 22/31] feat(log): add SimpleLogger for terminal logging without external backends --- src/experiment_logger/__init__.py | 2 + src/experiment_logger/simple_logger.py | 83 ++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 src/experiment_logger/simple_logger.py diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index 93b1de3..64d4be4 100644 --- a/src/experiment_logger/__init__.py +++ b/src/experiment_logger/__init__.py @@ -6,10 +6,12 @@ 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, get_logger +from experiment_logger.simple_logger import SimpleLogger from experiment_logger.wandb_utils import finish_wandb, init_wandb __all__ = [ "UnifiedLogger", + "SimpleLogger", "get_logger", "init_wandb", "finish_wandb", diff --git a/src/experiment_logger/simple_logger.py b/src/experiment_logger/simple_logger.py new file mode 100644 index 0000000..0ed1ab5 --- /dev/null +++ b/src/experiment_logger/simple_logger.py @@ -0,0 +1,83 @@ +"""Simple terminal logger for running without external backends. + +This is used for standalone package usage where WandB or TensorBoard are not desired. +It preserves the same API as UnifiedLogger but simply prints to stdout. +""" + +import logging +from typing import Any, Dict, Optional + + +class SimpleLogger: + """Simple logger that implements the UnifiedLogger interface via print statements.""" + + def __init__( + self, + run_name: str = "simple_run", + config: Optional[Dict[str, Any]] = None, + project_name: str = "none", + entity: Optional[str] = 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 {} + print(f"[INIT] SimpleLogger initialized for run: {run_name}") + + def set_level(self, level: int): + pass + + def log_non_interactive(self, msg: str, *args, **kwargs): + """In SimpleLogger, we just print everything as we assume interactive use.""" + self.info(msg, *args, **kwargs) + + def progress_bar(self, iterable=None, *args, **kwargs): + """Standard tqdm wrapper that falls back to range if tqdm is missing.""" + try: + import tqdm + + return tqdm.tqdm(iterable, *args, **kwargs) + except ImportError: + return iterable + + def info(self, msg: str, *args, **kwargs): + print(f"[INFO] {msg}") + + def warning(self, msg: str, *args, **kwargs): + print(f"[WARNING] {msg}") + + def error(self, msg: str, *args, **kwargs): + print(f"[ERROR] {msg}") + + def debug(self, msg: str, *args, **kwargs): + print(f"[DEBUG] {msg}") + + def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True): + step_str = f"Step {step}" if step is not None else "Log" + metric_str = ", ".join(f"{k}: {v}" for k, v in metrics.items()) + print(f"[{step_str}] {metric_str}") + + def save_checkpoint( + self, + params: Any, + step: int, + prefix: str = "checkpoint", + metadata: Optional[Dict[str, Any]] = None, + ): + print(f"[SAVE] Checkpoint '{prefix}' would be saved at step {step} (SimpleLogger: No-Op)") + + def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None): + print("[SAVE] Final model would be saved (SimpleLogger: No-Op)") + + def finish(self): + print(f"[FINISH] SimpleLogger finished for run: {self.run_name}") + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.finish() From e9b52e9e8fa2f174d0f6850adf061e21c77dc2ba Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 20:37:47 +0200 Subject: [PATCH 23/31] refactor: migrate configs from JSON to YAML --- configs/example.json | 8 ------- configs/example.yaml | 5 +++++ scripts/simulate.py | 4 ++-- .../environment/env_config.py | 12 ++++------- src/experiment_logger/unified_logger.py | 21 ++++++++++--------- 5 files changed, 22 insertions(+), 28 deletions(-) delete mode 100644 configs/example.json create mode 100644 configs/example.yaml diff --git a/configs/example.json b/configs/example.json deleted file mode 100644 index 8b646db..0000000 --- a/configs/example.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "morphology": { - "num_arms": 2, - "num_segments_per_arm": 4, - "use_p_control": true, - "use_torque_control": false - } -} \ No newline at end of file diff --git a/configs/example.yaml b/configs/example.yaml new file mode 100644 index 0000000..14bf3d7 --- /dev/null +++ b/configs/example.yaml @@ -0,0 +1,5 @@ +morphology: + num_arms: 2 + num_segments_per_arm: 4 + use_p_control: true + use_torque_control: false diff --git a/scripts/simulate.py b/scripts/simulate.py index 28062fa..957f6fa 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -10,7 +10,7 @@ from brittle_star_project import ( SimulationConfig, simulate_policy, ) -from brittle_star_project.environment import from_json +from brittle_star_project.environment import from_file from brittle_star_project.rl import RLModel # imports concrete models via rl.__init__ from brittle_star_project.rl.base import get_rl_model_registry @@ -44,7 +44,7 @@ def parse_args() -> argparse.Namespace: def main() -> None: args = parse_args() - morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json") + morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml") # ======= ENVIRONMENT SETUP ======= diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index a693286..78083e9 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -1,7 +1,6 @@ from __future__ import annotations from dataclasses import dataclass, field -import json from .env_types import Task @@ -51,14 +50,11 @@ class EnvConfig: def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]: - """Load configurations from a JSON or YAML file.""" - with open(path, "r") as f: - if path.endswith(".yaml") or path.endswith(".yml"): - import yaml + """Load configurations from a YAML file.""" + import yaml - config_dict = yaml.safe_load(f) - else: - config_dict = json.load(f) + with open(path, "r") as f: + config_dict = yaml.safe_load(f) morphology = MorphologyConfig(**config_dict.get("morphology", {})) arena = ArenaConfig(**config_dict.get("arena", {})) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 71380b0..40229c5 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -7,9 +7,9 @@ This logger ensures all experimental data is preserved by writing to: """ import datetime -import json import logging import subprocess +import yaml import sys import time from pathlib import Path @@ -99,7 +99,7 @@ class UnifiedLogger: self.metrics_dir = self.run_dir / "metrics" self.metrics_dir.mkdir(exist_ok=True) - self.config_file = self.run_dir / "config.json" + self.config_file = self.run_dir / "config.yaml" # Setup standard Python logging mirror self.text_log_file = self.run_dir / "run.log" @@ -196,7 +196,7 @@ class UnifiedLogger: """Save configuration to disk.""" try: with open(self.config_file, "w") as f: - json.dump(self.config, f, indent=2) + yaml.dump(self.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}") @@ -263,10 +263,10 @@ class UnifiedLogger: return try: - metrics_file = self.metrics_dir / "metrics.jsonl" + metrics_file = self.metrics_dir / "metrics.yaml" with open(metrics_file, "a") as f: for metric in self.metrics_buffer: - # Convert numpy/jax types to native Python types for JSON serialization + # Convert numpy/jax types to native Python types for YAML serialization serializable_metric = {} for k, v in metric.items(): if hasattr(v, "item"): # numpy/jax scalar @@ -275,7 +275,8 @@ class UnifiedLogger: serializable_metric[k] = v.tolist() else: serializable_metric[k] = v - f.write(json.dumps(serializable_metric) + "\n") + f.write("---\n") + yaml.dump(serializable_metric, f, default_flow_style=False) self.metrics_buffer.clear() except Exception as e: self.error(f"Error flushing metrics: {e}") @@ -298,9 +299,9 @@ class UnifiedLogger: # Save metadata if provided if metadata: - metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.json" + metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.yaml" with open(metadata_path, "w") as f: - json.dump(metadata, f, indent=2) + yaml.dump(metadata, f, default_flow_style=False) self.info(f"Checkpoint saved: {checkpoint_path}") @@ -334,9 +335,9 @@ class UnifiedLogger: f.write(flax.serialization.to_bytes(params)) if metadata: - metadata_path = self.run_dir / "final_model_metadata.json" + metadata_path = self.run_dir / "final_model_metadata.yaml" with open(metadata_path, "w") as f: - json.dump(metadata, f, indent=2) + yaml.dump(metadata, f, default_flow_style=False) self.info(f"Final model saved: {final_model_path}") From cd7169a75f0332720d61351a6f1590a05cc6794a Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 20:47:48 +0200 Subject: [PATCH 24/31] docs(log): extend how to use UnifiedLogger --- docs/DEVELOPMENT.md | 22 +++++++++++ src/experiment_logger/README.md | 67 +++++++++++++++++++++++++++++---- 2 files changed, 81 insertions(+), 8 deletions(-) diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index 6ec220b..5a01d26 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -59,3 +59,25 @@ Verify your setup by running the JAX initialization test: uv run pytest tests/test_jax_init.py ``` In the devcontainer, this will succeed on both CPU and GPU. A `GpuDevice` is expected if a GPU is detected and the `cuda` extra was installed. + +## Logging & Monitoring + +This project uses a unified logging system through the `experiment_logger` package. For a full API reference, see the [package README](../src/experiment_logger/README.md). + +### Quick Setup + +1. **Authorization**: Export your API key in your terminal to enable WandB synchronization: + ```bash + export WANDB_API_KEY=your_copied_api_key_here + ``` +2. **Toggle Tracking**: Use the `--track` flag in `scripts/train.py` to enable online sync. +3. **Local Monitoring**: All runs are recorded in the `runs/` directory. View scalars with TensorBoard: + ```bash + tensorboard --logdir runs/ + ``` + +### Environment Awareness + +The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job). It will automatically disable progress bars and switch to robust fallback modes (offline logging) to ensure your experiments never hang. + + diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md index ab6394d..ce31d69 100644 --- a/src/experiment_logger/README.md +++ b/src/experiment_logger/README.md @@ -1,17 +1,68 @@ # Experiment Logger -A lightweight logging framework supporting Weights & Biases, local JSON, and stdout. +A standardized, unified interface for logging experiments across multiple backends (WandB, TensorBoard, and Local Disk). -## Usage +This library is designed to be a standalone package that decouples the logging logic from the core training routines in the `brittle_star_project`. + +## Quick Start + +The recommended way to use the logger is through the `get_logger()` singleton: ```python -from experiment_logger import UnifiedLogger +from experiment_logger import get_logger -logger = UnifiedLogger(run_name="my_experiment", config={"lr": 0.001}) +logger = get_logger() +# Initialize at the start of your script (e.g., in train.py) +logger.init( + project_name="MyProject", + run_name="my_experiment_run", + base_dir="runs", + use_wandb=True +) -logger.log({"loss": 0.5}, step=1) -logger.save_checkpoint(params=model_params, step=1) -logger.finish() +# Log metrics (Scalar values, numpy scalars, or JAX types) +logger.log({"loss": 0.5, "accuracy": 0.98}, step=100) + +# Standard logging (Mirrored to disk and stdout) +logger.info("Training started") +logger.warning("Learning rate is very high") + +# Save checkpoints (Automatically synced to WandB as artifacts) +logger.save_checkpoint(params, step=5000) ``` -Logs and checkoints are saved in the `runs/` directory. If `track=True` (or `use_wandb=True`), everything is additionally synced to Weights & Biases. +## Logger Classes + +### `UnifiedLogger` + +The full suite for production training. It manages: +- **WandB**: Syncs metrics and uploads model checkpoints as artifacts. +- **TensorBoard**: Writes events for local visualization. +- **Local Disk**: Stores metrics in `metrics.yaml` and textual logs in `run.log`. + +### `SimpleLogger` + +A zero-dependency fallback that uses standard Python `print()` statements. Use this for standalone testing or minimal environments where you don't need persistent monitoring. + +```python +from experiment_logger import SimpleLogger +logger = SimpleLogger(run_name="test_run") +``` + +## API Features + +### `logger.progress_bar(iterable, **kwargs)` + +A smart wrapper around `tqdm` that automatically detects its environment. +- **Interactive Terminal**: Displays a normal progress bar. +- **Non-Interactive (HPC)**: Automatically disables the bar to prevent log file bloat in `slurm.out`. + +### `logger.log_non_interactive(msg: str)` + +Prints a message *only* when running in non-interactive environments. Useful for high-level progress tracking (e.g., "Epoch 5 Complete") without interactive noise. + +### `logger.save_checkpoint(params, step, prefix="checkpoint")` + +Saves model parameters using Flax serialization. +- **Local Location**: `runs//checkpoints/` +- **WandB Logic**: Automatically uploads the `.flax` file as a model artifact for lineage tracking. From 3c6eec2410f7836cbef144dbaf69151971dcd6a0 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 22:29:29 +0200 Subject: [PATCH 25/31] fix(hpc): typos, .env and pythonpath --- scripts/hpc/install.sh | 2 +- scripts/hpc/train.pbs | 14 +++++++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/scripts/hpc/install.sh b/scripts/hpc/install.sh index d1b69db..f88d081 100644 --- a/scripts/hpc/install.sh +++ b/scripts/hpc/install.sh @@ -19,7 +19,7 @@ if [ -n "$PBS_O_WORKDIR" ]; then cd "$PBS_O_WORKDIR" fi -mkdir "${PBS_O_WORKDIR}/runs" +mkdir -p "${PBS_O_WORKDIR}/runs" # Mirror configs to $VSC_DATA to avoid home quota limits (3GB) # vsc-venv manages environments relative to the requirements file diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs index 3d53193..e7d21f0 100644 --- a/scripts/hpc/train.pbs +++ b/scripts/hpc/train.pbs @@ -53,8 +53,20 @@ echo ">>> Starting BrittleStar training..." export MUJOCO_GL=egl export WANDB_DIR="$SCRATCH_RUNDIR" -python src/train.py \ +export PYTHONPATH="$PBS_O_WORKDIR/src:${PYTHONPATH:-}" + +if [ -f "$VSC_DATA/$PROJ_NAME/.env" ]; then + echo ">>> Sourcing API keys from .env..." + export $(grep -v '^#' "$VSC_DATA/$PROJ_NAME/.env" | xargs) +elif [ -f "$PBS_O_WORKDIR/.env" ]; then + echo ">>> Sourcing API keys from .env..." + export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs) +fi + +# TODO Once experiments get serious, change the config +python scripts/train.py \ --env-config-path configs/hpc/smoke_test.yaml \ + --hyperparameter-config-path configs/hpc/smoke_test.yaml \ --run-dir "$SCRATCH_RUNDIR" echo ">>> Staging out results to $DATA_RUNDIR..." From f6dc9c8e7fcfe7c16a7c6f48096a7edd3e18198e Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 23:18:20 +0200 Subject: [PATCH 26/31] feat: add build system, mypy and editable source --- pyproject.toml | 37 +++++++++++++++++++ src/__init__.py | 25 ------------- src/{ => brittle_star_project}/MLPs/mlps.py | 0 src/{ => brittle_star_project}/ppo.py | 0 .../trainers/PPOTrainer.py | 4 +- uv.lock | 2 +- 6 files changed, 40 insertions(+), 28 deletions(-) delete mode 100644 src/__init__.py rename src/{ => brittle_star_project}/MLPs/mlps.py (100%) rename src/{ => brittle_star_project}/ppo.py (100%) diff --git a/pyproject.toml b/pyproject.toml index b85ce0f..5687167 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,7 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + [project] name = "2026sel3-project" version = "0.1.0" @@ -41,3 +45,36 @@ dev = [ "pytest>=8.0.0", "ruff>=0.15.2", ] + +[tool.hatch.build.targets.wheel] +packages = ["src/brittle_star_project", "src/experiment_logger"] + +[tool.mypy] +mypy_path = "src" +check_untyped_defs = false +warn_return_any = false + +[[tool.mypy.overrides]] +module = [ + "jax.*", + "flax.*", + "wandb.*", + "torch.*", + "mujoco.*", + "mujoco_warp.*", + "optax.*", + "tyro.*", + "biorobot.*", + "gymnasium.*", + "matplotlib.*", + "mediapy.*", + "matplotlib.*", + "mediapy.*", + "pytest.*", + "tensorboard.*", + "tqdm.*", + "numpy.*", + "yaml.*", + "moojoco.*" +] +ignore_missing_imports = true diff --git a/src/__init__.py b/src/__init__.py deleted file mode 100644 index 35ecd10..0000000 --- a/src/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -from .brittle_star_project import ( - ArenaConfig, - Backend, - BrittleStarEnv, - BrittleStarEnvFactory, - EnvConfig, - MorphologyConfig, - Task, - simulate_policy, - SimulationConfig, - ControlPolicy, -) - -__all__ = [ - "ArenaConfig", - "Backend", - "BrittleStarEnv", - "BrittleStarEnvFactory", - "EnvConfig", - "MorphologyConfig", - "Task", - "simulate_policy", - "SimulationConfig", - "ControlPolicy", -] diff --git a/src/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py similarity index 100% rename from src/MLPs/mlps.py rename to src/brittle_star_project/MLPs/mlps.py diff --git a/src/ppo.py b/src/brittle_star_project/ppo.py similarity index 100% rename from src/ppo.py rename to src/brittle_star_project/ppo.py diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 1c807d4..bd97b9e 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -16,14 +16,14 @@ from experiment_logger import get_logger from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper -from MLPs.mlps import ( +from brittle_star_project.MLPs.mlps import ( Actor, AgentParams, GenericDenseLayersWithActivation, OneDenseLayerMLP, Storage, ) -from ppo import PPO +from brittle_star_project.ppo import PPO @jax.jit diff --git a/uv.lock b/uv.lock index 42321e1..01caa5d 100644 --- a/uv.lock +++ b/uv.lock @@ -11,7 +11,7 @@ resolution-markers = [ [[package]] name = "2026sel3-project" version = "0.1.0" -source = { virtual = "." } +source = { editable = "." } dependencies = [ { name = "biorobot" }, { name = "cleanrl" }, From 75c6d5bacda319a39a45f5aed3acbee49c3435b6 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 23:37:01 +0200 Subject: [PATCH 27/31] fix: mypy complaints --- scripts/train.py | 8 ++++---- src/brittle_star_project/dataclasses/PPOArgs.py | 3 --- src/brittle_star_project/render/renderer.py | 3 +++ src/brittle_star_project/trainers/PPOTrainer.py | 4 ++-- src/experiment_logger/README.md | 11 +++++++---- src/experiment_logger/config_utils.py | 14 +++++++------- src/experiment_logger/unified_logger.py | 6 +++--- 7 files changed, 26 insertions(+), 23 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 39b062d..944956b 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -10,7 +10,7 @@ from brittle_star_project.dataclasses import PPOArgs from brittle_star_project.trainers.PPOTrainer import PPOTrainer from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper -from experiment_logger.unified_logger import get_logger +from experiment_logger import UnifiedLogger from experiment_logger.config_utils import merge_config_with_cli, print_config @@ -37,7 +37,7 @@ def get_git_hash() -> str: return ( subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() ) - except subprocess.CalledProcessError | UnicodeDecodeError: + except (subprocess.CalledProcessError, UnicodeDecodeError): return "none" @@ -59,8 +59,8 @@ if __name__ == "__main__": os.makedirs(run_dir, exist_ok=True) # Initialize Global Logger - logger = get_logger() - logger.init( + logger = UnifiedLogger( + config=vars(args), project_name=args.wandb_project_name, # or default PPO-Modularity if missing run_name=run_name, base_dir=os.path.dirname(run_dir), diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 68265fa..036b44f 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/src/brittle_star_project/dataclasses/PPOArgs.py @@ -22,9 +22,6 @@ class PPOArgs: # the directory to save the experiment results run_dir: str | None = None - # how often to save checkpoints (0 to disable) - checkpoint_frequency: int = 0 - # seed of the experiment seed: int = 1 diff --git a/src/brittle_star_project/render/renderer.py b/src/brittle_star_project/render/renderer.py index 4d97ef1..91e669c 100644 --- a/src/brittle_star_project/render/renderer.py +++ b/src/brittle_star_project/render/renderer.py @@ -36,6 +36,9 @@ def simulate_policy( import mujoco.viewer + if state is None: + raise ValueError("A valid environment state must be provided.") + model = state.mj_model data = state.mj_data diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index bd97b9e..d4297de 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -312,14 +312,14 @@ class PPOTrainer: def _init_episode_stats(self) -> EpisodeStatistics: self.logger.info("[EPISODE STATS]: Initializing episode stats...") - return EpisodeStatistics( + return EpisodeStatistics( # type: ignore[call-arg] episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32), episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32), returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), ) - def _rollout(self, env_state, next_obs, next_done) -> tuple[Storage, ...]: + def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]: return self._rollout_jit( self.agent_state, self.episode_stats, diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md index ce31d69..4e13b6b 100644 --- a/src/experiment_logger/README.md +++ b/src/experiment_logger/README.md @@ -9,17 +9,20 @@ This library is designed to be a standalone package that decouples the logging l The recommended way to use the logger is through the `get_logger()` singleton: ```python -from experiment_logger import get_logger +from experiment_logger import UnifiedLogger, get_logger -logger = get_logger() # Initialize at the start of your script (e.g., in train.py) -logger.init( - project_name="MyProject", +logger = UnifiedLogger( run_name="my_experiment_run", + config={"learning_rate": 3e-4}, + project_name="MyProject", base_dir="runs", use_wandb=True ) +# In other files, retrieve the initialized singleton: +# logger = get_logger() + # Log metrics (Scalar values, numpy scalars, or JAX types) logger.log({"loss": 0.5, "accuracy": 0.98}, step=100) diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index 7a92c0e..4c5b79f 100644 --- a/src/experiment_logger/config_utils.py +++ b/src/experiment_logger/config_utils.py @@ -44,10 +44,10 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: raise ValueError(f"{cls} is not a dataclass") # Get field names and types - field_map = {f.name: f for f in fields(cls)} + field_map = {f.name: f for f in fields(cls)} # type: ignore # Filter config to only include valid fields - filtered_config = {} + filtered_config: Dict[str, Any] = {} for key, value in config_dict.items(): if key in field_map: field = field_map[key] @@ -64,7 +64,7 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: 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 + filtered_config[key] = field.type(value) if value is not None else None # type: ignore except (ValueError, TypeError) as e: log.warning(f"Could not convert {key}={value} to {field.type}: {e}") filtered_config[key] = value @@ -74,7 +74,7 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: return cls(**filtered_config) -def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T: +def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T: """Merge YAML config with CLI arguments, with CLI taking precedence. Args: @@ -107,16 +107,16 @@ def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T: # 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)} + default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)} # type: ignore # Parse CLI args cli_instance = tyro.cli(config_class) - cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)} + cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)} # type: ignore # Merge configs: YAML as base, CLI overrides non-default values final_config = {} - for field in fields(config_class): + for field in fields(config_class): # type: ignore field_name = field.name default_value = default_dict[field_name] yaml_value = yaml_config.get(field_name, default_value) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 40229c5..051fe22 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -224,7 +224,7 @@ class UnifiedLogger: self._log_to_stdout(metrics_with_metadata) # Log to WandB - if self.wandb_available: + if self.wandb_run is not None: try: self.wandb_run.log(metrics, step=step, commit=commit) except Exception as e: @@ -306,7 +306,7 @@ class UnifiedLogger: self.info(f"Checkpoint saved: {checkpoint_path}") # Log to WandB as artifact - if self.wandb_available: + if self.wandb_run is not None: try: import wandb @@ -342,7 +342,7 @@ class UnifiedLogger: self.info(f"Final model saved: {final_model_path}") # Log to WandB - if self.wandb_available: + if self.wandb_run is not None: try: import wandb From fe10faa344a72efd41dc4019048821749808ab7c Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 9 Apr 2026 00:04:36 +0200 Subject: [PATCH 28/31] test(log): wandb integration config file --- configs/hpc/wandb_test.yaml | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 configs/hpc/wandb_test.yaml diff --git a/configs/hpc/wandb_test.yaml b/configs/hpc/wandb_test.yaml new file mode 100644 index 0000000..bf2229d --- /dev/null +++ b/configs/hpc/wandb_test.yaml @@ -0,0 +1,11 @@ +# Configuration to verify WandB online tracking +exp_name: "hpc_wandb_verification" +seed: 42 +track: true # Enabled for testing WandB +wandb_project_name: "PPO-Modularity" +wandb_entity: "SEL3-2026-Groep-4" + +num_envs: 128 +total_timesteps: 50000 # Short run for quick verification +num_steps: 128 +cuda: true From 2f31df412b841ada48477c433dfe9e264f01b345 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 9 Apr 2026 00:14:28 +0200 Subject: [PATCH 29/31] fix(log): avoid duplicates in logs --- src/experiment_logger/unified_logger.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 051fe22..5cfb3f1 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -105,6 +105,7 @@ class UnifiedLogger: self.text_log_file = self.run_dir / "run.log" self._text_logger = logging.getLogger(f"UnifiedLogger_{self.run_name}") self._text_logger.setLevel(log_level) + self._text_logger.propagate = False # Avoid duplicate handlers if re-instantiated if not self._text_logger.handlers: From 849581fcfd2ee53838679b55494cbd7305715cba Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 9 Apr 2026 00:33:56 +0200 Subject: [PATCH 30/31] fix: linting --- scripts/train.py | 2 -- src/brittle_star_project/trainers/PPOTrainer.py | 1 - src/experiment_logger/wandb_utils.py | 3 ++- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 944956b..492c887 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -2,8 +2,6 @@ import subprocess import time import torch -import tyro -import yaml import os from brittle_star_project.dataclasses import PPOArgs diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index d4297de..8c238f1 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -5,7 +5,6 @@ from dataclasses import asdict, dataclass from functools import partial from typing import Any -import flax import jax import jax.numpy as jnp import numpy as np diff --git a/src/experiment_logger/wandb_utils.py b/src/experiment_logger/wandb_utils.py index 2c162fb..302d308 100644 --- a/src/experiment_logger/wandb_utils.py +++ b/src/experiment_logger/wandb_utils.py @@ -52,7 +52,8 @@ def init_wandb( if not has_key and not is_interactive and os.environ.get("WANDB_MODE") != "offline": logger.warning( - "WANDB_API_KEY not found and environment is non-interactive. Switching to offline mode." + "WANDB_API_KEY not found and environment is non-interactive. " + "Switching to offline mode." ) sync_path = f"runs/{name}" if name else "runs" logger.warning(f"WandB is offline. Use 'wandb sync {sync_path}' to upload logs later.") From 2d43f5e6422acd42c781a75392233fbb66e12bb2 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 9 Apr 2026 15:02:02 +0200 Subject: [PATCH 31/31] Apply suggestions from code review Co-authored-by: RobinMeersman <77965843+RobinMeersman@users.noreply.github.com> Co-authored-by: Tibo De Peuter --- .gitignore | 3 --- README.md | 6 +++--- configs/README.md | 6 +++--- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/.gitignore b/.gitignore index 79b1594..b0567ce 100644 --- a/.gitignore +++ b/.gitignore @@ -3,9 +3,6 @@ artifacts/* runs/* wandb/ -# Experiment tracking -wandb/ - # Python-generated files __pycache__/ *.py[oc] diff --git a/README.md b/README.md index 19c994d..e6e5e78 100644 --- a/README.md +++ b/README.md @@ -34,19 +34,19 @@ uv sync --frozen example command: ```bash -uv run python src/train.py +uv run python scripts/train.py ``` Or use a custom config file: ```bash -uv run python src/train.py --config configs/my_experiment.yaml +uv run python scripts/train.py --config configs/my_experiment.yaml ``` Override specific parameters: ```bash -uv run python src/train.py --learning-rate 0.001 --num-envs 32 --track +uv run python scripts/train.py --learning-rate 0.001 --num-envs 32 --track ``` ### Logging diff --git a/configs/README.md b/configs/README.md index c1662fb..58c6aed 100644 --- a/configs/README.md +++ b/configs/README.md @@ -4,15 +4,15 @@ This directory contains configuration files for training experiments. ## Usage -Use `--config` with `src/train.py` to run an experiment: +Use `--config` with `scripts/train.py` to run an experiment: ```bash -python src/train.py --config configs/default_ppo.yaml +python scripts/train.py --config configs/default_ppo.yaml ``` You can overriding settings via CLI: ```bash -python src/train.py --config configs/default_ppo.yaml --learning-rate 0.001 +python scripts/train.py --config configs/default_ppo.yaml --learning-rate 0.001 ``` ## Available Configurations