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/README.md b/README.md index 085f3d5..e6e5e78 100644 --- a/README.md +++ b/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. diff --git a/configs/README.md b/configs/README.md new file mode 100644 index 0000000..58c6aed --- /dev/null +++ b/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. diff --git a/configs/default_ppo.yaml b/configs/default_ppo.yaml new file mode 100644 index 0000000..1b06c3d --- /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: "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 diff --git a/configs/dev_test.yaml b/configs/dev_test.yaml new file mode 100644 index 0000000..b64d330 --- /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: "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 \ No newline at end of file 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/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 diff --git a/configs/personal_template.yaml b/configs/personal_template.yaml new file mode 100644 index 0000000..67caef7 --- /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: "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 \ No newline at end of file diff --git a/configs/production_training.yaml b/configs/production_training.yaml index 215b2be..6dce29b 100644 --- a/configs/production_training.yaml +++ b/configs/production_training.yaml @@ -1,24 +1,42 @@ -# Full PPO training config for Brittle Star (HPC Production) -exp_name: "production_training" -seed: 1 -track: true -capture_video: true -save_model: true -checkpoint_frequency: 100 # not yet implemented in train.py but here for future use +# Production Training Configuration +# +# Full-scale training configuration for production runs +# with wandb logging enabled. -# Scaling for HPC (using A100 GPU slices) -num_envs: 128 -total_timesteps: 10000000 -num_steps: 128 +# Experiment settings +exp_name: "brittle_star_production_training" +seed: 42 + +# 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 \ No newline at end of file diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md index e694bb3..bd58706 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 (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. 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/pyproject.toml b/pyproject.toml index 171bae1..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" @@ -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 diff --git a/ruff.toml b/ruff.toml index db1bcd3..ad24132 100644 --- a/ruff.toml +++ b/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 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..." 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/scripts/train.py b/scripts/train.py index 9ee6da6..492c887 100644 --- a/scripts/train.py +++ b/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: - 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 @@ -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 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/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 3f04c16..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 @@ -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 diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 7c86d51..b143d9c 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/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) 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/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/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/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 7df57f0..8c238f1 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/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): - 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 +263,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,18 +308,17 @@ 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( + 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( - "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, iteration: int) -> tuple: + if iteration == 1: + self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}") ( 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: - print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True) + if iteration == 1: + 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: - print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True) + 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 log and not is_tty and iteration == 1: - print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True) + 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() @@ -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): - 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...") - - if not is_tty: - print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True) + self.logger.info("[TRAIN]: Resetting environment...") + 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 log and not is_tty: - print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True) + self.logger.log_non_interactive(f"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, - ) + 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 log and 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)) - print( - 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, - ) + 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/README.md b/src/experiment_logger/README.md new file mode 100644 index 0000000..4e13b6b --- /dev/null +++ b/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//checkpoints/` +- **WandB Logic**: Automatically uploads the `.flax` file as a model artifact for lineage tracking. diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py new file mode 100644 index 0000000..64d4be4 --- /dev/null +++ b/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" diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py new file mode 100644 index 0000000..4c5b79f --- /dev/null +++ b/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}") 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() diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py new file mode 100644 index 0000000..5cfb3f1 --- /dev/null +++ b/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() diff --git a/src/experiment_logger/wandb_utils.py b/src/experiment_logger/wandb_utils.py new file mode 100644 index 0000000..302d308 --- /dev/null +++ b/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}") 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) diff --git a/uv.lock b/uv.lock index fd45cfc..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" }, @@ -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" },