Merge pull request #20 from SELab-3-2026/feat/wandb-logging
feat: Experiment logging setup
This commit is contained in:
GitHub
commit
70facf51cf
35 changed files with 1303 additions and 220 deletions
12
.env.example
Normal file
12
.env.example
Normal file
|
|
@ -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
|
||||
48
README.md
48
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.
|
||||
|
|
|
|||
23
configs/README.md
Normal file
23
configs/README.md
Normal file
|
|
@ -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.
|
||||
48
configs/default_ppo.yaml
Normal file
48
configs/default_ppo.yaml
Normal file
|
|
@ -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
|
||||
42
configs/dev_test.yaml
Normal file
42
configs/dev_test.yaml
Normal file
|
|
@ -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
|
||||
|
|
@ -1,8 +0,0 @@
|
|||
{
|
||||
"morphology": {
|
||||
"num_arms": 2,
|
||||
"num_segments_per_arm": 4,
|
||||
"use_p_control": true,
|
||||
"use_torque_control": false
|
||||
}
|
||||
}
|
||||
5
configs/example.yaml
Normal file
5
configs/example.yaml
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
morphology:
|
||||
num_arms: 2
|
||||
num_segments_per_arm: 4
|
||||
use_p_control: true
|
||||
use_torque_control: false
|
||||
11
configs/hpc/wandb_test.yaml
Normal file
11
configs/hpc/wandb_test.yaml
Normal file
|
|
@ -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
|
||||
40
configs/personal_template.yaml
Normal file
40
configs/personal_template.yaml
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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..."
|
||||
|
|
|
|||
|
|
@ -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 =======
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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", {}))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
71
src/experiment_logger/README.md
Normal file
71
src/experiment_logger/README.md
Normal file
|
|
@ -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.
|
||||
21
src/experiment_logger/__init__.py
Normal file
21
src/experiment_logger/__init__.py
Normal file
|
|
@ -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"
|
||||
147
src/experiment_logger/config_utils.py
Normal file
147
src/experiment_logger/config_utils.py
Normal file
|
|
@ -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}")
|
||||
83
src/experiment_logger/simple_logger.py
Normal file
83
src/experiment_logger/simple_logger.py
Normal file
|
|
@ -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()
|
||||
385
src/experiment_logger/unified_logger.py
Normal file
385
src/experiment_logger/unified_logger.py
Normal file
|
|
@ -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()
|
||||
91
src/experiment_logger/wandb_utils.py
Normal file
91
src/experiment_logger/wandb_utils.py
Normal file
|
|
@ -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}")
|
||||
38
tests/test_config.py
Normal file
38
tests/test_config.py
Normal file
|
|
@ -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)
|
||||
4
uv.lock
generated
4
uv.lock
generated
|
|
@ -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" },
|
||||
|
|
|
|||
Reference in a new issue