Merge pull request #20 from SELab-3-2026/feat/wandb-logging

feat: Experiment logging setup

.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

README.md

@@ -1,8 +1,8 @@
 # Brittle Star
 
-## Usage
+## Quick Start
 
-### UV
+### Installation
 
 To set up the UV module, you can run the following command:
 
@@ -10,12 +10,54 @@ To set up the UV module, you can run the following command:
 uv sync --frozen
 ```
 
+### 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
+
 example command:
 
 ```bash
-uv run src/train.py --model_name my_model --epochs 50 --batch_size 32
+uv run python scripts/train.py
 ```
 
+Or use a custom config file:
+
+```bash
+uv run python scripts/train.py --config configs/my_experiment.yaml
+```
+
+Override specific parameters:
+
+```bash
+uv run python scripts/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.
+
 ## HPC
 
 See **[docs/HPC.md](docs/HPC.md)** for the full guide, including environment setup, cluster selection, interactive debugging, and job submission.

configs/README.md

@@ -0,0 +1,23 @@
+# Configuration Files
+
+This directory contains configuration files for training experiments.
+
+## Usage
+
+Use `--config` with `scripts/train.py` to run an experiment:
+
+```bash
+python scripts/train.py --config configs/default_ppo.yaml
+```
+
+You can overriding settings via CLI:
+```bash
+python scripts/train.py --config configs/default_ppo.yaml --learning-rate 0.001
+```
+
+## Available Configurations
+
+- `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.

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: "SEL3-2026-Groep-4"  # 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

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: "SEL3-2026-Groep-4"  # ⚠️  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

configs/example.json

@@ -1,8 +0,0 @@
-{
-  "morphology": {
-    "num_arms": 2,
-    "num_segments_per_arm": 4,
-    "use_p_control": true,
-    "use_torque_control": false
-  }
-}

configs/example.yaml

@@ -0,0 +1,5 @@
+morphology:
+  num_arms: 2
+  num_segments_per_arm: 4
+  use_p_control: true
+  use_torque_control: false

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

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: "SEL3-2026-Groep-4"  # ⚠️ 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

configs/production_training.yaml

@@ -1,24 +1,42 @@
-# Full PPO training config for Brittle Star (HPC Production)
+# Production Training Configuration
-exp_name: "production_training"
+#
-seed: 1
+# Full-scale training configuration for production runs
-track: true
+# with wandb logging enabled.
-capture_video: true
-save_model: true
-checkpoint_frequency: 100  # not yet implemented in train.py but here for future use
 
-# Scaling for HPC (using A100 GPU slices)
+# Experiment settings
-num_envs: 128
+exp_name: "brittle_star_production_training"
-total_timesteps: 10000000
+seed: 42
-num_steps: 128
+
+# Tracking
+track: true
+capture_video: false
+wandb_project_name: "PPO-Modularity"
+wandb_entity: "SEL3-2026-Groep-4"
+
+# Model saving
+save_model: true
+checkpoint_frequency: 100  # Save checkpoint every 100 iterations
+
+# Environment settings
+num_envs: 512
+
+# Training hyperparameters
+total_timesteps: 50000000
+num_steps: 256
 num_minibatches: 4
 update_epochs: 4
 
-# Algorithm
 learning_rate: 2.5e-4
 anneal_lr: true
 gamma: 0.99
 gae_lambda: 0.95
 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

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 (Google standard). This is enforced via `uv` using **ruff** and pre-commit hooks.
+
+## 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.

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.
+
+

pyproject.toml

@@ -1,3 +1,7 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
 [project]
 name = "2026sel3-project"
 version = "0.1.0"
@@ -21,6 +25,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",
     "torch>=2.4.0",
@@ -40,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

ruff.toml

@@ -350,7 +350,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

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

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

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

scripts/train.py

@@ -2,14 +2,15 @@ import subprocess
 import time
 
 import torch
-import tyro
-import yaml
 import os
 
 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 import UnifiedLogger
+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 +18,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:
+def parse_args() -> PPOArgs:
-    temp_args = tyro.cli(PPOArgs)
+    import argparse
 
-    if temp_args.hyperparameter_config_path is not None:
+    # Use argparse to reliably extract just the config path without swallowing --help
-        if log:
+    parser = argparse.ArgumentParser(add_help=False)
-            print(f"Loading hyperparameter config from {temp_args.hyperparameter_config_path}")
+    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:
+    args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path)
-            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
     return args
 
 
@@ -47,7 +35,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"
 
 
@@ -60,6 +48,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 +56,17 @@ if __name__ == "__main__":
 
     os.makedirs(run_dir, exist_ok=True)
 
+    # Initialize Global Logger
+    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),
+        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

src/__init__.py

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

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
 
@@ -41,7 +38,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
@@ -49,6 +46,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
 

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -35,6 +35,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
@@ -52,6 +59,7 @@ class BrittleStarJaxEnvWrapper:
         return self._env.observation_space
 
     def reset(self, seed: int = 0):
+        self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
         self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
         env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
         return self._vectorized_reset(rng=env_rngs)

src/brittle_star_project/environment/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."""
+    """Load configurations from a YAML file."""
-    with open(path, "r") as f:
+    import yaml
-        if path.endswith(".yaml") or path.endswith(".yml"):
-            import yaml
 
-            config_dict = yaml.safe_load(f)
+    with open(path, "r") as f:
-        else:
+        config_dict = yaml.safe_load(f)
-            config_dict = json.load(f)
 
         morphology = MorphologyConfig(**config_dict.get("morphology", {}))
         arena = ArenaConfig(**config_dict.get("arena", {}))

src/brittle_star_project/environment/factory.py

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

src/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
 

src/brittle_star_project/trainers/PPOTrainer.py

@@ -1,30 +1,28 @@
 import datetime
 import random
-import sys
 import time
 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
 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
-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
@@ -207,7 +205,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 +245,14 @@ class PPOTrainer:
 
         self._init_random()
 
-    def _init_random(self, log: bool = True):
+    def _init_random(self):
-        if log:
+        self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}")
-            print(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):
+    def _init_agent(self):
-        if log:
+        self.logger.info("[AGENT]: Initializing agent...")
-            print("[AGENT]: Initializing agent...")
 
         sensor = GenericDenseLayersWithActivation()
         feature_extractor = GenericDenseLayersWithActivation()
@@ -267,9 +263,8 @@ class PPOTrainer:
         # messenger = OneDenseLayerMLP()
         return sensor, feature_extractor, actor, critic
 
-    def _init_agent_state(self, log: bool = True) -> TrainState:
+    def _init_agent_state(self) -> TrainState:
-        if log:
+        self.logger.info("[AGENT STATE]: Initializing agent state...")
-            print("[AGENT STATE]: Initializing agent state...")
 
         self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split(
             self.key, 5
@@ -313,18 +308,17 @@ class PPOTrainer:
             ),
         )
 
-    def _init_episode_stats(self, log: bool = True) -> EpisodeStatistics:
+    def _init_episode_stats(self) -> EpisodeStatistics:
-        if log:
+        self.logger.info("[EPISODE STATS]: Initializing episode stats...")
-            print("[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,
@@ -350,39 +344,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(
+    def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple:
-            "charts/avg_episodic_return", loss_info.avg_episodic_return, global_step
+        if iteration == 1:
-        )
+            self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}")
-        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)
 
         (
             self.agent_state,
@@ -394,20 +378,20 @@ class PPOTrainer:
             next_env_state,
         ) = self._rollout(env_state, next_obs, next_done)
 
-        if log and not is_tty and iteration == 1:
+        if iteration == 1:
-            print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True)
+            self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}")
 
         storage = self._compute_gae(storage, next_obs, next_done)
 
-        if log and not is_tty and iteration == 1:
+        if iteration == 1:
-            print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True)
+            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 log and not is_tty and iteration == 1:
+        if iteration == 1:
-            print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True)
+            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()
@@ -429,105 +413,64 @@ class PPOTrainer:
 
     def _close(self):
         self.env.close()
-        self.writer.close()
 
-    def _save_model(self, model_path: str, log: bool = True):
+    def _save_model(self, model_path: str):
-        if log:
+        self.logger.info("[SAVE]: Saving the final model...")
-            print(f"[SAVE]: Saving the model to: {model_path}...")
 
-        with open(model_path, "wb") as f:
+        params = [
-            f.write(
+            vars(self.args),
-                flax.serialization.to_bytes(
+            [
-                    [
+                self.agent_state.params["sensor_params"],
-                        vars(self.args),
+                self.agent_state.params["actor_params"],
-                        [
+                self.agent_state.params["critic_params"],
-                            self.agent_state.params["sensor_params"],
+                self.agent_state.params["feature_extractor_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:
+        self.logger.info(f"running name: {self.run_name}")
-            print(f"running name: {self.run_name}")
 
-        is_tty = sys.stdout.isatty()
+        self.logger.info("[TRAIN]: Resetting environment...")
-        if log:
+        self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
-            print("[TRAIN]: Resetting environment...")
-
-            if not is_tty:
-                print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True)
 
         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:
+        self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
-            print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True)
 
         global_step = 0
         start_time = time.time()
 
-        if self.args.track:
+        iter_bar = self.logger.progress_bar(range(1, self.args.num_iterations + 1))
-            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,
-        )
         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 log and not is_tty:
+            sps = int(global_step / (time.time() - start_time))
-                sps = int(global_step / (time.time() - start_time))
+            remaining_steps = self.args.total_timesteps - global_step
-                remaining_steps = self.args.total_timesteps - global_step
+            eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
-                eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
+            eta_str = str(datetime.timedelta(seconds=eta_seconds))
-                eta_str = str(datetime.timedelta(seconds=eta_seconds))
 
-                print(
+            self.logger.log_non_interactive(
-                    f"Iteration {iteration}/{self.args.num_iterations} | "
+                f"Iteration {iteration}/{self.args.num_iterations} | "
-                    f"Step {global_step}/{self.args.total_timesteps} | "
+                f"Step {global_step}/{self.args.total_timesteps} | "
-                    f"SPS {sps} | "
+                f"SPS {sps} | "
-                    f"Return {loss_info.avg_episodic_return:.4f} | "
+                f"Return {loss_info.avg_episodic_return:.4f} | "
-                    f"ETA {eta_str}",
+                f"ETA {eta_str}"
-                    flush=True,
+            )
-                )
 
         if self.args.save_model:
             model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model"

src/experiment_logger/README.md

@@ -0,0 +1,71 @@
+# Experiment Logger
+
+A standardized, unified interface for logging experiments across multiple backends (WandB, TensorBoard, and Local Disk).
+
+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, get_logger
+
+# Initialize at the start of your script (e.g., in train.py)
+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)
+
+# 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)
+```
+
+## 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/<run_name>/checkpoints/`
+- **WandB Logic**: Automatically uploads the `.flax` file as a model artifact for lineage tracking.

src/experiment_logger/__init__.py

@@ -0,0 +1,21 @@
+"""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.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",
+    "load_yaml_config",
+    "merge_config_with_cli",
+]
+__version__ = "0.1.0"

src/experiment_logger/config_utils.py

@@ -0,0 +1,147 @@
+"""Configuration utilities for loading YAML configs and merging with CLI args."""
+
+import os
+import sys
+from typing import Dict, Any, Type, TypeVar
+import yaml
+from dataclasses import fields, is_dataclass
+
+from experiment_logger.unified_logger import get_logger
+
+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:
+        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)}  # type: ignore
+
+    # Filter config to only include valid fields
+    filtered_config: Dict[str, Any] = {}
+    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 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  # type: ignore
+            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 = 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)}  # 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)}  # type: ignore
+
+    # Merge configs: YAML as base, CLI overrides non-default values
+    final_config = {}
+
+    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)
+        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}")

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

src/experiment_logger/unified_logger.py

@@ -0,0 +1,385 @@
+"""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, run.log)
+3. stdout (for real-time monitoring)
+"""
+
+import datetime
+import logging
+import subprocess
+import yaml
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import flax
+import jax.numpy as jnp
+import numpy as np
+
+from experiment_logger.wandb_utils import finish_wandb, init_wandb
+
+# 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_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(
+            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:
+    """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,
+        log_level: int = logging.INFO,
+        _set_as_global: 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
+            _set_as_global: Internal flag to override the global singleton
+        """
+        self.run_name = run_name
+        self.config = config
+        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
+        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.yaml"
+
+        # 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(log_level)
+        self._text_logger.propagate = False
+
+        # 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()
+
+        # 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)
+
+        # Initialize metrics storage
+        self.metrics_buffer: List[Dict[str, Any]] = []
+        self.step_counter = 0
+
+        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 set_level(self, level: int):
+        """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)
+
+    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."""
+        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:
+                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}")
+
+    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_run is not None:
+            try:
+                self.wandb_run.log(metrics, step=step, commit=commit)
+            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)
+
+        # 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"]
+        )
+        self.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.yaml"
+            with open(metrics_file, "a") as f:
+                for metric in self.metrics_buffer:
+                    # 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
+                            serializable_metric[k] = v.item()
+                        elif isinstance(v, (np.ndarray, jnp.ndarray)):
+                            serializable_metric[k] = v.tolist()
+                        else:
+                            serializable_metric[k] = v
+                    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}")
+
+    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."""
+        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.yaml"
+                with open(metadata_path, "w") as f:
+                    yaml.dump(metadata, f, default_flow_style=False)
+
+            self.info(f"Checkpoint saved: {checkpoint_path}")
+
+            # Log to WandB as artifact
+            if self.wandb_run is not None:
+                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)
+                    self.info("Checkpoint uploaded to WandB")
+                except Exception as e:
+                    self.warning(f"Could not upload checkpoint to WandB: {e}")
+
+        except Exception as 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."""
+        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.yaml"
+                with open(metadata_path, "w") as f:
+                    yaml.dump(metadata, f, default_flow_style=False)
+
+            self.info(f"Final model saved: {final_model_path}")
+
+            # Log to WandB
+            if self.wandb_run is not None:
+                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:
+                    self.warning(f"Could not upload final model to WandB: {e}")
+
+        except Exception as e:
+            self.error(f"Error saving final model: {e}")
+
+    def finish(self):
+        """Finalize logging and cleanup."""
+        # 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
+        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()

src/experiment_logger/wandb_utils.py

@@ -0,0 +1,91 @@
+"""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
+        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,
+            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}")

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)

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" },
@@ -28,6 +28,7 @@ dependencies = [
     { name = "protobuf" },
     { name = "pyopengl" },
     { name = "pyopengl-accelerate" },
+    { name = "pyyaml" },
     { name = "torch" },
     { name = "tyro" },
     { name = "wandb" },
@@ -67,6 +68,7 @@ requires-dist = [
     { name = "protobuf", specifier = ">=5.0.0" },
     { name = "pyopengl", specifier = ">=3.1.10" },
     { name = "pyopengl-accelerate", specifier = ">=3.1.10" },
+    { name = "pyyaml", specifier = ">=6.0" },
     { name = "tensorboard", marker = "extra == 'analysis'" },
     { name = "torch", specifier = ">=2.4.0" },
     { name = "tyro", specifier = ">=1.0.10" },