diff --git a/.gitignore b/.gitignore index b0567ce..83cfc3f 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,9 @@ artifacts/* runs/* wandb/ +outputs/ +multirun/ +metrics/ # Python-generated files __pycache__/ diff --git a/configs/.gitkeep b/configs/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/configs/README.md b/configs/README.md index 58c6aed..91208f2 100644 --- a/configs/README.md +++ b/configs/README.md @@ -1,23 +1,54 @@ -# Configuration Files +# Brittle Star Configuration System -This directory contains configuration files for training experiments. +This project uses **Hydra** for a modular, hierarchical, and strictly-typed configuration system. -## Usage +## Core Concepts -Use `--config` with `scripts/train.py` to run an experiment: +1. **Composition over Inheritance**: Instead of one giant config file, the configuration is composed of small, domain-specific modules (PPO settings, architecture, morphology, etc.). +2. **Strict Typing**: Every configuration is validated against a Python dataclass schema (`ConfigStore`). Misspelled keys throw a `ConfigAttributeError` immediately. +3. **CLI Swapping**: You can swap entire modules or override individual values from the command line without touching code. +## Directory Structure + +- `main_config.yaml`: The root entry point defining the default composition. +- `experiment/`: High-level experiment settings (seed, device). +- `logging/`: WandB and checkpointing configuration. +- `ppo/`: PPO training hyperparameters. +- `architecture/`: Polymorphic network architectures (centralized vs. decentralized). +- `morphology/`: Physical robot definitions (number of segments, amputations). +- `arena/`: Environment physics and visual settings. +- `environment/`: Task-specific settings (Directed Locomotion, Light Escape). + +## Common Commands + +### Local Debugging +Run a quick test with minimal iterations: ```bash -python scripts/train.py --config configs/default_ppo.yaml +python scripts/train.py experiment=dev_test ppo=fast ``` -You can overriding settings via CLI: +### Swapping Architectures or Morphologies +Test a decentralized controller on a 3-arm robot: ```bash -python scripts/train.py --config configs/default_ppo.yaml --learning-rate 0.001 +python scripts/train.py architecture=decentralized morphology=3_arms ``` -## Available Configurations +### HPC Production +Run stable PPO with WandB enabled (HPC submission scripts handle the `hydra.run.dir` redirection): +```bash +python scripts/train.py ppo=stable logging=wandb_enabled +``` -- `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. +### Dry-Run Validation +Check if your configuration is valid without starting the simulation: +```bash +python scripts/train.py --cfg job +``` + +## Developer Notes + +- **Adding a new group**: Create a subdirectory in `configs/` and register the new dataclass in `src/brittle_star_project/configs/register_configs.py`. +- **Typo Catching**: If you see a `ConfigAttributeError`, check for typos in your YAML keys or CLI overrides. +- **Output Redirection**: We use `experiment.base_run_dir` to configure where logs and models are stored (defaults to `runs/`). + - To change it locally: `python scripts/train.py experiment.base_run_dir=/path/to/custom/dir` + - On HPC, ensure this points to a fast scratch storage. diff --git a/configs/architecture/centralized.yaml b/configs/architecture/centralized.yaml new file mode 100644 index 0000000..bb4bc0f --- /dev/null +++ b/configs/architecture/centralized.yaml @@ -0,0 +1,22 @@ +# Centralized Actor-Critic Architecture +# Baseline configuration with a single global sensor and motor. + +# Default values are defined in CentralizedConfig dataclass. +# Use this configuration for standard PPO experiments. + +name: "centralized" +sensor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +motor: + hidden_dims: [] + activation: "tanh" + +feature_extractor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +critic: + hidden_dims: [] + activation: "tanh" diff --git a/configs/architecture/decentralized.yaml b/configs/architecture/decentralized.yaml new file mode 100644 index 0000000..4b9fffb --- /dev/null +++ b/configs/architecture/decentralized.yaml @@ -0,0 +1,32 @@ +# Decentralized Actor Architecture (NerveNet-MLP variant) +# Multi-agent/distributed configuration using local sensors, propagators, and motors. + +# Default values are defined in DecentralizedConfig dataclass. +# Use this configuration for decentralized execution experiments. + +name: "decentralized" +sensor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +propagator: + hidden_dims: [300, 300, 300] + activation: "tanh" + +motor: + hidden_dims: [] + activation: "tanh" + +feature_extractor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +critic: + hidden_dims: [] + activation: "tanh" + +# Synchronous message-passing rounds per control step +message_passing_steps: 1 + +# Connectivity topology (e.g., ring, fully_connected) +topology_type: "ring" diff --git a/configs/arena/default.yaml b/configs/arena/default.yaml new file mode 100644 index 0000000..c42ec3c --- /dev/null +++ b/configs/arena/default.yaml @@ -0,0 +1,8 @@ +# Default Arena Configuration +# Base aquarium environment settings. + +size: [10.0, 5.0] +sand_ground_color: true +attach_target: true +wall_height: 1.5 +wall_thickness: 0.1 diff --git a/configs/default_ppo.yaml b/configs/default_ppo.yaml deleted file mode 100644 index 1b06c3d..0000000 --- a/configs/default_ppo.yaml +++ /dev/null @@ -1,48 +0,0 @@ -# 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 deleted file mode 100644 index b64d330..0000000 --- a/configs/dev_test.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# 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/environment/directed_locomotion.yaml b/configs/environment/directed_locomotion.yaml new file mode 100644 index 0000000..63ad520 --- /dev/null +++ b/configs/environment/directed_locomotion.yaml @@ -0,0 +1,12 @@ +# Directed Locomotion Environment +# Baseline task setting. + +task: DIRECTED_LOCOMOTION +simulation_time: 5.0 +num_physics_steps_per_control_step: 10 +time_scale: 2 +camera_ids: [0, 1] +render_size: [480, 640] +joint_randomization_noise_scale: 0.0 +target_distance: 3.0 +light_perlin_noise_scale: 0 diff --git a/configs/environment/light_escape.yaml b/configs/environment/light_escape.yaml new file mode 100644 index 0000000..c79cfb9 --- /dev/null +++ b/configs/environment/light_escape.yaml @@ -0,0 +1,12 @@ +# Light Escape Environment +# Advanced task requiring movement away from light source. + +task: LIGHT_ESCAPE +simulation_time: 5.0 +num_physics_steps_per_control_step: 10 +time_scale: 2 +camera_ids: [0, 1] +render_size: [480, 640] +joint_randomization_noise_scale: 0.0 +target_distance: 3.0 +light_perlin_noise_scale: 200 # Must be integer factor of 200 diff --git a/configs/example.yaml b/configs/example.yaml deleted file mode 100644 index 14bf3d7..0000000 --- a/configs/example.yaml +++ /dev/null @@ -1,5 +0,0 @@ -morphology: - num_arms: 2 - num_segments_per_arm: 4 - use_p_control: true - use_torque_control: false diff --git a/configs/experiment/base.yaml b/configs/experiment/base.yaml new file mode 100644 index 0000000..9f31d27 --- /dev/null +++ b/configs/experiment/base.yaml @@ -0,0 +1,7 @@ +# Base Experiment Configuration +# Default values align with ExperimentConfig dataclass. + +exp_name: "brittle_star_ppo" +seed: 1 +torch_deterministic: true +cuda: true diff --git a/configs/experiment/dev_test.yaml b/configs/experiment/dev_test.yaml new file mode 100644 index 0000000..941a1e9 --- /dev/null +++ b/configs/experiment/dev_test.yaml @@ -0,0 +1,7 @@ +# Testing Experiment Configuration +# Quick experiment for local development/testing. + +exp_name: "dev_test_brittle_star" +seed: 42 +torch_deterministic: true +cuda: true diff --git a/configs/experiment/hpc_smoke_test.yaml b/configs/experiment/hpc_smoke_test.yaml new file mode 100644 index 0000000..4264801 --- /dev/null +++ b/configs/experiment/hpc_smoke_test.yaml @@ -0,0 +1,7 @@ +# HPC Smoke Test Configuration +# Uses minimal settings but simulates HPC environment. + +exp_name: "hpc_smoke_test" +seed: 123 +torch_deterministic: true +cuda: true diff --git a/configs/hpc/smoke_test.yaml b/configs/hpc/smoke_test.yaml deleted file mode 100644 index 1dd2fc0..0000000 --- a/configs/hpc/smoke_test.yaml +++ /dev/null @@ -1,11 +0,0 @@ -# Minimal config to verify HPC setup is functional. -# Run with: python scripts/train.py --config-path configs/hpc/smoke_test.yaml -exp_name: "hpc_smoke_test" -seed: 0 -track: false # Test WandB integration -capture_video: false # No rendering for smoke test -save_model: true # Test the end-of-training save routine -num_envs: 512 -total_timesteps: 65536 -num_steps: 128 -cuda: true diff --git a/configs/hpc/wandb_expand.yaml b/configs/hpc/wandb_expand.yaml deleted file mode 100644 index 432160f..0000000 --- a/configs/hpc/wandb_expand.yaml +++ /dev/null @@ -1,10 +0,0 @@ -exp_name: "explained_var_fun_more_steps" -seed: 42 -track: true -wandb_project_name: "LET-THERE-BE-MORE-LOGGING" -wandb_entity: "SEL3-2026-Groep-4" - -num_envs: 16 -num_steps: 256 -total_timesteps: 50000 -cuda: true \ No newline at end of file diff --git a/configs/hpc/wandb_test.yaml b/configs/hpc/wandb_test.yaml deleted file mode 100644 index bf2229d..0000000 --- a/configs/hpc/wandb_test.yaml +++ /dev/null @@ -1,11 +0,0 @@ -# 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/logging/default.yaml b/configs/logging/default.yaml new file mode 100644 index 0000000..4a0cb92 --- /dev/null +++ b/configs/logging/default.yaml @@ -0,0 +1,11 @@ +# Default Logging Configuration +# Offline local-only setup (WandB disabled). + +track: false +wandb_project_name: "PPO-Modularity" +wandb_entity: "SEL3-2026-Groep-4" +capture_video: false +save_model: true +checkpoint_frequency: 100 +upload_model: false +hf_entity: "" diff --git a/configs/logging/wandb_enabled.yaml b/configs/logging/wandb_enabled.yaml new file mode 100644 index 0000000..2e82781 --- /dev/null +++ b/configs/logging/wandb_enabled.yaml @@ -0,0 +1,11 @@ +# WandB Enabled Logging Configuration +# For production/cloud experiments with weights synced. + +track: true +wandb_project_name: "PPO-Modularity" +wandb_entity: "SEL3-2026-Groep-4" +capture_video: false +save_model: true +checkpoint_frequency: 100 +upload_model: false +hf_entity: "" diff --git a/configs/main_config.yaml b/configs/main_config.yaml new file mode 100644 index 0000000..745755a --- /dev/null +++ b/configs/main_config.yaml @@ -0,0 +1,21 @@ +# Brittle Star Project - Main Configuration +# This file defines the default composition of the hierarchical configuration. +# Sub-configs are loaded from the relative directories. + +defaults: + - brittle_star_config + - experiment: base + - logging: default + - ppo: default + - architecture: centralized + - morphology: 5_arms_full + - arena: default + - environment: directed_locomotion + - simulation: default + - _self_ + +hydra: + job: + chdir: True + run: + dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/configs/morphology/3_arms.yaml b/configs/morphology/3_arms.yaml new file mode 100644 index 0000000..ebc1665 --- /dev/null +++ b/configs/morphology/3_arms.yaml @@ -0,0 +1,6 @@ +# 3 Arms Morphology Configuration +# Symmetric amputation (arms 1 and 3 removed). + +segments_per_arm: [4, 0, 4, 0, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/5_arms_full.yaml b/configs/morphology/5_arms_full.yaml new file mode 100644 index 0000000..408b7c9 --- /dev/null +++ b/configs/morphology/5_arms_full.yaml @@ -0,0 +1,6 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 4, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/partial_amputation.yaml b/configs/morphology/partial_amputation.yaml new file mode 100644 index 0000000..8003035 --- /dev/null +++ b/configs/morphology/partial_amputation.yaml @@ -0,0 +1,6 @@ +# Partial Amputation Configuration +# Random partial amputation for robustness testing. + +segments_per_arm: [4, 2, 4, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/personal_template.yaml b/configs/personal_template.yaml deleted file mode 100644 index 67caef7..0000000 --- a/configs/personal_template.yaml +++ /dev/null @@ -1,40 +0,0 @@ -# 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/ppo/debug.yaml b/configs/ppo/debug.yaml new file mode 100644 index 0000000..58dbd3a --- /dev/null +++ b/configs/ppo/debug.yaml @@ -0,0 +1,16 @@ +learning_rate: 0.0003 +total_timesteps: 409600 +num_envs: 32 +num_steps: 32 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 32 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.005 +vf_coef: 1.0 +max_grad_norm: 0.5 +target_kl: null diff --git a/configs/ppo/default.yaml b/configs/ppo/default.yaml new file mode 100644 index 0000000..50107b4 --- /dev/null +++ b/configs/ppo/default.yaml @@ -0,0 +1,19 @@ +# Default PPO Configuration +# Standard hyperparams from original codebase. + +learning_rate: 0.00025 +total_timesteps: 10000000 +num_envs: 100 +num_steps: 128 +anneal_lr: true +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 diff --git a/configs/ppo/fast.yaml b/configs/ppo/fast.yaml new file mode 100644 index 0000000..5001ef2 --- /dev/null +++ b/configs/ppo/fast.yaml @@ -0,0 +1,19 @@ +# Fast PPO Configuration +# Lower timestep count for quick iterations/testing. + +learning_rate: 0.0005 +total_timesteps: 500000 +num_envs: 8 +num_steps: 128 +anneal_lr: true +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 diff --git a/configs/ppo/smoke_test.yaml b/configs/ppo/smoke_test.yaml new file mode 100644 index 0000000..40e4c07 --- /dev/null +++ b/configs/ppo/smoke_test.yaml @@ -0,0 +1,19 @@ +# Fast PPO Configuration +# Lower timestep count for quick iterations/testing. + +learning_rate: 0.0005 +total_timesteps: 65536 +num_envs: 512 +num_steps: 128 +anneal_lr: true +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 diff --git a/configs/ppo/stable.yaml b/configs/ppo/stable.yaml new file mode 100644 index 0000000..35cade4 --- /dev/null +++ b/configs/ppo/stable.yaml @@ -0,0 +1,19 @@ +# Stable PPO Configuration +# Standard hyperparams with lower LR and larger batch. + +learning_rate: 0.0001 +total_timesteps: 10000000 +num_envs: 100 +num_steps: 256 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 8 +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 diff --git a/configs/production_training.yaml b/configs/production_training.yaml deleted file mode 100644 index 6dce29b..0000000 --- a/configs/production_training.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Production Training Configuration -# -# Full-scale training configuration for production runs -# with wandb logging enabled. - -# 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 - -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/configs/simulation/default.yaml b/configs/simulation/default.yaml new file mode 100644 index 0000000..00c4e44 --- /dev/null +++ b/configs/simulation/default.yaml @@ -0,0 +1,11 @@ +# Default Simulation Settings +# These values are used by scripts/simulate.py + +# Path to the trained model (optional) +model_path: null + +# Type of model to use if no path is provided (e.g., random) +model_type: "random" + +# Execution backend (MJX or BRAX) +backend: "MJX" diff --git a/env/hpc/modules.txt b/env/hpc/modules.txt index 4c4a1d7..2ea2c6e 100644 --- a/env/hpc/modules.txt +++ b/env/hpc/modules.txt @@ -1,3 +1,4 @@ GCCcore/13.3.0 Python/3.12.3-GCCcore-13.3.0 FFmpeg/7.0.2-GCCcore-13.3.0 +Hydra/1.3.2-GCCcore-13.3.0 diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md new file mode 100644 index 0000000..92f23a1 --- /dev/null +++ b/experiments/debug-experiment-10042026/used_variables.md @@ -0,0 +1,107 @@ +## Default envconfig +task: Task = Task.DIRECTED_LOCOMOTION +simulation_time: float = 500.0 +num_physics_steps_per_control_step: int = 10 +time_scale: int = 2 +camera_ids: list[int] = field(default_factory=lambda: [0, 1]) +render_size: tuple[int, int] = (480, 640) +joint_randomization_noise_scale: float = 0.0 +target_distance: float = 3.0 +light_perlin_noise_scale: int = 0 + + +## Default ppoargs +seed: int = 1 +torch_deterministic: bool = True +cuda: bool = True +track: bool = False +checkpoint_frequency: int = 100 +learning_rate: float = 2.5e-4 +anneal_lr: bool = True +gamma: float = 0.99 +gae_lambda: float = 0.95 +update_epochs: int = 4 +norm_adv: bool = True +clip_vloss: bool = True +max_grad_norm: float = 0.5 +target_kl: float | None = None +batch_size: int = 0 +minibatch_size: int = 0 +num_iterations: int = 0 + +## Used config file: (hpc/debug.yaml) +exp_name: "debug-experiment" +seed: 42 +track: true +wandb_project_name: "Let's-find-that-bug" +wandb_entity: "SEL3-2026-Groep-4" +run_dir: "/data/gent/465/vsc46589" +num_envs: 32 +num_steps: 32 +num_minibatches: 32 +total_timesteps: 409600 +num_arms: 2 +cuda: true + +ent_coef: 0.005 +vf_coef: 1.0 +clip_coef: 0.2 + +anneal_lr: true +learning_rate: 0.0003 + +## Arena config: +size: tuple[float, float] = (10.0, 5.0) +sand_ground_color: bool = True +attach_target: bool = True +wall_height: float = 1.5 +wall_thickness: float = 0.1 + +## Morphology: +num_segments_per_arm: int = 4 +use_p_control: bool = True +use_torque_control: bool = False + +## MLPs: +### Sensor & Feature_extractor: +Both with 3 layers of 300 neurons per layer. + +class GenericDenseLayersWithActivation(nn.Module): + layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64]) + activation: Callable = nn.tanh + + @nn.compact + def __call__(self, x): + for size in self.layer_sizes: + x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x) + x = self.activation(x) + return x + +### Actor: +class Actor(nn.Module): + action_dim: int + @nn.compact + def __call__(self, x): + mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x) + log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,)) + return mean, log_std + +### Critic: +class OneDenseLayerMLP(nn.Module): + @nn.compact + def __call__(self, x): + return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x) + +### Observations: +_ALLOWED_OBS_KEYS = { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + "disk_position", + "disk_rotation", + "disk_linear_velocity", + "disk_angular_velocity", + "unit_xy_direction_to_target", + "xy_distance_to_target", +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 5687167..607e704 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "pyopengl>=3.1.10", "pyopengl-accelerate>=3.1.10", "pyyaml>=6.0", - "tyro>=1.0.10", + "hydra-core>=1.3.2", "wandb==0.24.2", "torch>=2.4.0", ] diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs index b9c8519..60ba4d8 100644 --- a/scripts/hpc/train.pbs +++ b/scripts/hpc/train.pbs @@ -63,11 +63,11 @@ elif [ -f "$PBS_O_WORKDIR/.env" ]; then export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs) fi -# TODO Once experiments get serious, change the config +# Run training using Hydra overrides python scripts/train.py \ - --env-config-path configs/hpc/wandb_expand.yaml \ - --hyperparameter-config-path configs/hpc/wandb_expand.yaml \ - --run-dir "$SCRATCH_RUNDIR" + hydra.run.dir="$SCRATCH_RUNDIR" \ + ppo=stable \ + logging=wandb_enabled echo ">>> Staging out results to $DATA_RUNDIR..." cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/" diff --git a/scripts/simulate.py b/scripts/simulate.py index b6d170f..e9a4ef3 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -1,30 +1,54 @@ +"""Simulate a trained policy in the MuJoCo viewer. + +Uses Hydra to load the same BrittleStarConfig that was used during training. +Override settings via CLI, e.g.: + python scripts/simulate.py morphology=3_arms +""" + from __future__ import annotations -import argparse import itertools import time from pathlib import Path from typing import Any + import flax +import hydra import jax import jax.numpy as jnp import numpy as np +from omegaconf import DictConfig, OmegaConf -from brittle_star_project import ( - Backend, -) -from brittle_star_project.environment import ArenaConfig, EnvConfig, MorphologyConfig, from_file +from brittle_star_project import BrittleStarEnv, BrittleStarEnvFactory +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs +_ALLOWED_OBS_KEYS = { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + "disk_position", + "disk_rotation", + "disk_linear_velocity", + "disk_angular_velocity", + "unit_xy_direction_to_target", + "xy_distance_to_target", +} -def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray: +def _transform_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray: """Flatten the env's observation dict into a 1D vector. - concatenates values in the dict's iteration order and skips empty arrays. + Matches training behavior: + - only includes keys in _ALLOWED_OBS_KEYS + - iterates keys in sorted order for stable layout + - skips empty arrays """ - parts: list[jnp.ndarray] = [] - for v in obs_dict.values(): - arr = jnp.asarray(v) + for key in sorted(obs_dict.keys()): + if key not in _ALLOWED_OBS_KEYS: + continue + arr = jnp.asarray(obs_dict[key]) if arr.size == 0: continue parts.append(arr.reshape((-1,))) @@ -45,7 +69,9 @@ class CleanRLPPOPolicy: ) -> None: from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation - self._sensor = GenericDenseLayersWithActivation() + hidden_dim = int(sensor_params["params"]["Dense_0"]["kernel"].shape[1]) + + self._sensor = GenericDenseLayersWithActivation(layer_sizes=[hidden_dim, hidden_dim]) self._actor = Actor(action_dim=action_dim) self._sensor_apply = jax.jit(self._sensor.apply) self._actor_apply = jax.jit(self._actor.apply) @@ -77,28 +103,29 @@ class CleanRLPPOPolicy: def _parse_checkpoint(restored_obj: Any) -> tuple[Any, Any, Any, Any, Any]: """Extract checkpoint parts. - Returns (args_dict, sensor_params, actor_params, critic_params, + Returns (config_dict, sensor_params, actor_params, critic_params, feature_extractor_params). - `PPOTrainer` saves: - flax.serialization.to_bytes( - [vars(args), [sensor, actor, critic, feature_extractor]] - ) + PPOTrainer saves: + flax.serialization.to_bytes([ + config_dict, + [sensor_params, actor_params, critic_params, feature_extractor_params], + ]) - `msgpack_restore()` may restore lists as dicts keyed by string - indices ("0", "1", ...), so we accept both shapes. + msgpack_restore() may restore lists as dicts keyed by string indices + ("0", "1", ...), so we accept both shapes. """ - args_part: Any | None = None + cfg_part: Any | None = None params_part: Any = restored_obj if isinstance(restored_obj, (list, tuple)) and len(restored_obj) >= 2: - args_part = restored_obj[0] + cfg_part = restored_obj[0] params_part = restored_obj[1] elif _looks_like_indexed_dict(restored_obj) and ( "0" in restored_obj or "1" in restored_obj ): - args_part = restored_obj.get("0", restored_obj.get(0)) + cfg_part = restored_obj.get("0", restored_obj.get(0)) params_part = restored_obj.get("1", restored_obj.get(1)) if _looks_like_indexed_dict(params_part): @@ -109,7 +136,7 @@ class CleanRLPPOPolicy: if sensor_params is None or actor_params is None: raise ValueError("Missing required params in checkpoint") return ( - args_part, + cfg_part, sensor_params, actor_params, critic_params, @@ -122,7 +149,7 @@ class CleanRLPPOPolicy: critic_params = params_part[2] if len(params_part) >= 3 else None feature_extractor_params = params_part[3] if len(params_part) >= 4 else None return ( - args_part, + cfg_part, sensor_params, actor_params, critic_params, @@ -131,14 +158,13 @@ class CleanRLPPOPolicy: raise ValueError( f"Unexpected checkpoint structure in {path}. " - "Expected [args_dict, [sensor_params, actor_params, critic_params, " - "feature_extractor_params]] " - "or an equivalent dict-indexed variant." + "Expected [config_dict, [sensor_params, actor_params, critic_params, " + "feature_extractor_params]] or an equivalent dict-indexed variant." ) payload = path.read_bytes() restored = flax.serialization.msgpack_restore(payload) - _args_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = ( + _cfg_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = ( _parse_checkpoint(restored) ) @@ -149,7 +175,7 @@ class CleanRLPPOPolicy: ) def act(self, *, observations: dict[str, Any]) -> np.ndarray: - obs = _flatten_obs_dict(observations) + obs = _transform_obs_dict(observations) hidden = self._sensor_apply(self._params["sensor_params"], obs) mean, _log_std = self._actor_apply(self._params["actor_params"], hidden) @@ -172,27 +198,26 @@ def _target_reached(*, state: Any) -> bool: def _rollout_one_episode_headless( *, - env: Any, + env: BrittleStarEnv, policy: CleanRLPPOPolicy, seed: int, max_steps: int, ) -> tuple[float, int, bool, float | None]: - """Run one rollout up to `max_steps`. + """Run one rollout up to max_steps. Returns (return, length, reached_target, final_xy_dist). + + Note: In the MJC backend, the raw env reward can be 0.0; we compute a simple + progress reward based on xy_distance_to_target. """ + state = env.reset(seed=seed) ep_return = 0.0 - observations = _get_observations(state) prev_dist = _get_xy_distance_to_target(observations) reached_target = _target_reached(state=state) - # NOTE: In the MJC backend, `state.reward` is always 0.0. - # To get a meaningful return, we compute a simple progress reward: - # r_t = d_{t-1} - d_t - # where d is `xy_distance_to_target`. steps = 0 for _ in range(int(max_steps)): action = policy.act(observations=observations) @@ -220,7 +245,7 @@ def _rollout_one_episode_headless( def _run_one_episode_viewer( *, - env: Any, + env: BrittleStarEnv, policy: CleanRLPPOPolicy, seed: int, state: Any, @@ -232,15 +257,14 @@ def _run_one_episode_viewer( model = state.mj_model data = state.mj_data - seed = int(seed) + _ = int(seed) episode_return = 0.0 observations = _get_observations(state) prev_dist = _get_xy_distance_to_target(observations) reached_target = _target_reached(state=state) steps = 0 - # Use the viewer as a context manager to avoid GLX teardown races - # (e.g. GLXBadDrawable from X_GLXSwapBuffers after a window is destroyed). + # Use the viewer as a context manager to avoid GLX teardown races. with mujoco.viewer.launch_passive(model, data) as viewer: step_iter = range(int(max_steps)) if max_steps is not None else itertools.count() for _step_idx in step_iter: @@ -248,7 +272,7 @@ def _run_one_episode_viewer( break step_start = time.time() - action = policy.act(observations=observations) + action = policy.act(observations=observations or {}) if model.nu > 0 and action.shape != (int(model.nu),): raise ValueError( f"Policy returned action shape {action.shape}, expected ({int(model.nu)},)" @@ -287,119 +311,86 @@ def _run_one_episode_viewer( ) -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser( - description="Run a trained policy for exactly one episode (viewer or headless)." - ) - p.add_argument( - "--config-path", - type=str, - default=None, - help=( - "Path to an environment YAML config (morphology/arena/env). " - "If omitted, uses the environment defaults. " - "Relative paths are resolved from the repository root." - ), - ) - p.add_argument( - "--model", - type=str, - required=True, - help=("Path to the Flax checkpoint saved by scripts/train.py (final_model.flax)."), - ) - p.add_argument( - "--headless", - action="store_true", - help="Run without the MuJoCo viewer (still exactly one episode).", - ) - p.add_argument( - "--max-steps", - type=int, - default=None, - help=( - "Number of control steps to run. " - "In --headless mode this is required and acts as a fixed horizon. " - "In viewer mode the default is infinite (run until window closed or target reached)." - ), - ) - p.add_argument( - "--backend", - choices=[b.value for b in Backend], - default=Backend.MJC.value, - ) - p.add_argument("--seed", type=int, default=0) - return p.parse_args() +def _infer_checkpoint_obs_dim(policy: CleanRLPPOPolicy) -> int | None: + """Best-effort read of the first Dense kernel input dim (obs dim).""" + + try: + kernel = policy._params["sensor_params"]["params"]["Dense_0"]["kernel"] + return int(getattr(kernel, "shape")[0]) + except Exception: + return None -def main() -> None: - from brittle_star_project.environment import ( - BrittleStarEnv, - BrittleStarEnvFactory, +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig) -> None: + # Convert DictConfig to structured dataclass, ensuring the root schema is applied. + config: BrittleStarConfig = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) ) - args = parse_args() + backend = config.simulation.backend + seed = int(config.experiment.seed) - if args.config_path is None: - morphology_cfg = MorphologyConfig() - arena_cfg = ArenaConfig() - env_cfg = EnvConfig() - else: - repo_root = Path(__file__).resolve().parents[1] - config_path = Path(args.config_path) - if not config_path.is_absolute(): - config_path = repo_root / config_path - morphology_cfg, arena_cfg, env_cfg = from_file(str(config_path)) + model_path_str = config.simulation.model_path + if model_path_str is None: + raise ValueError( + "simulation.model_path must be set to a .flax checkpoint (e.g. final_model.flax)" + ) + + # Hydra chdir changes CWD; resolve relative paths relative to the invocation. + model_path = Path(hydra.utils.to_absolute_path(model_path_str)) + if model_path.suffix != ".flax": + raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.") # ======= ENVIRONMENT SETUP ======= - - backend = Backend(args.backend) - factory = BrittleStarEnvFactory() - raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg) + raw_env = factory.create_environment( + backend, + config.morphology, + config.arena, + config.environment, + ) env = BrittleStarEnv( raw_env, backend=backend, - config=env_cfg, - morphology_config=morphology_cfg, + config=config.environment, + morphology_config=config.morphology, ) - seed_for_env = int(args.seed) if args.seed is not None else 0 - state = env.reset(seed=seed_for_env) + state0 = env.reset(seed=seed) # ======= MODEL SETUP ======= + nu = int(state0.mj_model.nu) + policy = CleanRLPPOPolicy.load(model_path, action_dim=nu) - # Extract the number of actuators (nu) from the environment's model, so we can pass it to the - # policy/model. - nu = int(state.mj_model.nu) - - model_path = Path(args.model) - if model_path.suffix != ".flax": + # Helpful early failure when configs don't match the checkpoint. + observations0 = _get_observations(state0) + env_obs_dim = int(_transform_obs_dict(observations0 or {}).shape[0]) + ckpt_obs_dim = _infer_checkpoint_obs_dim(policy) + if ckpt_obs_dim is not None and ckpt_obs_dim != env_obs_dim: raise ValueError( - f"Expected the training artifact '.flax', got '{model_path.name}'." + "Checkpoint/env mismatch: " + f"checkpoint expects obs_dim={ckpt_obs_dim}, env provides obs_dim={env_obs_dim}. " + "Use the same Hydra config (morphology/arena/environment) " + "that was used during training." ) - policy = CleanRLPPOPolicy.load( - model_path, - action_dim=nu, - ) - - default_seed = seed_for_env - # ======= SIMULATION ======= + headless = bool(config.simulation.headless) + max_steps = config.simulation.max_steps - if args.headless: - if args.max_steps is None: - raise ValueError("--max-steps is required in --headless mode") - max_steps = int(args.max_steps) - if max_steps <= 0: - raise ValueError("--max-steps must be > 0") + if headless: + if max_steps is None: + raise ValueError("simulation.max_steps is required when simulation.headless=true") + max_steps_i = int(max_steps) + if max_steps_i <= 0: + raise ValueError("simulation.max_steps must be > 0") - ep_seed = int(args.seed) if args.seed is not None else default_seed ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless( env=env, policy=policy, - seed=ep_seed, - max_steps=max_steps, + seed=seed, + max_steps=max_steps_i, ) final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}" print( @@ -408,27 +399,29 @@ def main() -> None: f"target_reached={reached_target}, final_xy_dist={final_dist_str}" ) else: - max_steps: int | None - if args.max_steps is None: - max_steps = None + if max_steps is not None: + max_steps_i = int(max_steps) + if max_steps_i <= 0: + raise ValueError("simulation.max_steps must be > 0") + max_steps_val: int | None = max_steps_i else: - max_steps = int(args.max_steps) - if max_steps <= 0: - raise ValueError("--max-steps must be > 0") + max_steps_val = None + + model_dt = float(state0.mj_model.opt.timestep) + control_dt = model_dt * float(config.environment.num_physics_steps_per_control_step) - model_dt = float(state.mj_model.opt.timestep) - control_dt = model_dt * float(env_cfg.num_physics_steps_per_control_step) _run_one_episode_viewer( env=env, policy=policy, - seed=int(args.seed) if args.seed is not None else default_seed, - state=state, + seed=seed, + state=state0, control_dt=control_dt, - max_steps=max_steps, + max_steps=max_steps_val, ) env.close() if __name__ == "__main__": + register_configs() main() diff --git a/scripts/train.py b/scripts/train.py index 492c887..409f050 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1,75 +1,60 @@ -import subprocess -import time - -import torch import os +import torch +import hydra +from omegaconf import DictConfig, OmegaConf -from brittle_star_project.dataclasses import PPOArgs +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs 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 +from experiment_logger import init_logger, get_logger -def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper: - if config_path is None: - return BrittleStarJaxEnvWrapper.default(num_envs=num_envs) - return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs) +def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: + """Create the environment using the structured configuration.""" + return BrittleStarJaxEnvWrapper( + morphology=cfg.morphology, + arena=cfg.arena, + env_config=cfg.environment, + num_envs=cfg.ppo.num_envs, + ) -def parse_args() -> PPOArgs: - import argparse +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig): + # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly. + config: BrittleStarConfig = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) + ) - # 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() + # 2. Setup run metadata + # Hydra changes CWD to the output directory by default. + run_dir = os.getcwd() + run_name = os.path.basename(run_dir) - args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path) - return args + # 3. Initialize Logger + cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True) + init_logger( + run_name=run_name, + config=cfg_dict, + project_name=config.logging.wandb_project_name, + entity=config.logging.wandb_entity, + base_dir=os.path.dirname(run_dir), + use_wandb=config.logging.track, + ) + logger = get_logger() + logger.info(f"Hydra-initialized run: {run_name}") + logger.info(f"Output directory: {run_dir}") + # 4. Setup Environment and Torch + env = make_env(config) + torch.backends.cudnn.deterministic = config.experiment.torch_deterministic -def get_git_hash() -> str: - try: - return ( - subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() - ) - except (subprocess.CalledProcessError, UnicodeDecodeError): - return "none" + # 5. Train - pass structured config directly + ppo_trainer = PPOTrainer(config, env, run_dir, run_name) + ppo_trainer.train() if __name__ == "__main__": - args = parse_args() - - args.batch_size = args.num_envs * args.num_steps - args.minibatch_size = args.batch_size // args.num_minibatches - args.num_iterations = args.total_timesteps // args.batch_size - - 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: - run_dir = args.run_dir - - 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 - - ppo_trainer = PPOTrainer(args, env, run_dir, run_name) - ppo_trainer.train() + register_configs() + main() diff --git a/src/brittle_star_project/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py index 6abb540..5e2deb5 100644 --- a/src/brittle_star_project/MLPs/mlps.py +++ b/src/brittle_star_project/MLPs/mlps.py @@ -58,6 +58,10 @@ class Storage: returns: jnp.array rewards: jnp.array + raw_actions: jnp.ndarray = None # before clipping + means: jnp.ndarray = None # policy mean + stds: jnp.ndarray = None # policy std + def replace(self, **kwargs) -> "Storage": fs = fields(self) return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs}) diff --git a/src/brittle_star_project/configs/config_architecture.py b/src/brittle_star_project/configs/config_architecture.py new file mode 100644 index 0000000..e8cf5a5 --- /dev/null +++ b/src/brittle_star_project/configs/config_architecture.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass, field +from typing import List, Optional + + +@dataclass +class LayerConfig: + hidden_dims: List[int] = field(default_factory=lambda: [64, 64]) + activation: str = "tanh" + + +@dataclass +class ArchitectureConfig: + """Base class for actor-critic network configurations. + + Both centralized and decentralized architectures share a centralized critic + composed of a feature extractor followed by a shallow output layer. + + See docs/design/actor-critic.md for the full design rationale. + """ + + name: str = "base" + + # Actor pipeline + sensor: Optional[LayerConfig] = None + propagator: Optional[LayerConfig] = None + motor: Optional[LayerConfig] = None + + # Critic pipeline + feature_extractor: Optional[LayerConfig] = None + critic: Optional[LayerConfig] = None + + # Decentralized + message_passing_steps: Optional[int] = None + topology_type: Optional[str] = None # Supported values: "ring", "fully_connected" + + +@dataclass +class CentralizedConfig(ArchitectureConfig): + """Centralized actor-critic architecture (baseline). + + The actor is a single global policy composed of a sensor (input network) + and a motor (output network). The sensor receives the full concatenated + global observation; the motor projects the hidden state to all joint actions. + + See docs/design/actor-critic.md for the full design rationale. + """ + + name: str = "centralized" + + +@dataclass +class DecentralizedConfig(ArchitectureConfig): + """Decentralized actor architecture (NerveNet-MLP variant). + + Each node runs a local sensor, exchanges messages with neighbours via a + propagator for a fixed number of steps, and then a local motor produces + the joint offset for that node only. + + The critic remains centralized (shared with the base class): it receives the + full concatenated global observation and outputs a single scalar. + + See docs/design/actor-critic.md and docs/design/communication.md for the + full design rationale. + """ + + name: str = "decentralized" diff --git a/src/brittle_star_project/configs/config_experiment.py b/src/brittle_star_project/configs/config_experiment.py new file mode 100644 index 0000000..7409811 --- /dev/null +++ b/src/brittle_star_project/configs/config_experiment.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass + + +@dataclass +class ExperimentConfig: + exp_name: str = "brittle_star_ppo" + seed: int = 1 + torch_deterministic: bool = True + cuda: bool = True + debug_sanity: bool = False + base_run_dir: str = "runs" diff --git a/src/brittle_star_project/configs/config_ppo.py b/src/brittle_star_project/configs/config_ppo.py new file mode 100644 index 0000000..d0a29bf --- /dev/null +++ b/src/brittle_star_project/configs/config_ppo.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class PPOConfig: + learning_rate: float = 2.5e-4 + total_timesteps: int = 10000000 + num_envs: int = 100 + num_steps: int = 128 + anneal_lr: bool = True + gamma: float = 0.99 + gae_lambda: float = 0.95 + num_minibatches: int = 4 + update_epochs: int = 4 + norm_adv: bool = True + clip_coef: float = 0.1 + clip_vloss: bool = True + ent_coef: float = 0.01 + vf_coef: float = 0.5 + max_grad_norm: float = 0.5 + target_kl: Optional[float] = None diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py new file mode 100644 index 0000000..56747ae --- /dev/null +++ b/src/brittle_star_project/configs/config_simulation.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass +from typing import Optional +from brittle_star_project.environment.env_types import Backend + + +@dataclass +class SimulationSettings: + """Settings for the simulation script.""" + + model_path: Optional[str] = None + model_type: str = "random" + backend: Backend = Backend.MJX diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py new file mode 100644 index 0000000..5a937a7 --- /dev/null +++ b/src/brittle_star_project/configs/main_config.py @@ -0,0 +1,28 @@ +from dataclasses import dataclass, field + +from experiment_logger.config_logger import LoggingConfig +from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_ppo import PPOConfig +from brittle_star_project.configs.config_architecture import ArchitectureConfig +from brittle_star_project.configs.config_simulation import SimulationSettings +from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig + + +@dataclass +class BrittleStarConfig: + """Root configuration for a brittle star training run. + + Composed of strictly separated sub-configs. Each sub-config can be swapped + independently via CLI or a different YAML file. See configs/README.md. + """ + + experiment: ExperimentConfig = field(default_factory=ExperimentConfig) + logging: LoggingConfig = field(default_factory=LoggingConfig) + ppo: PPOConfig = field(default_factory=PPOConfig) + # This field is polymorphic; defaults to the base class to allow subclasses + # (CentralizedConfig, DecentralizedConfig) to be merged in via Hydra. + architecture: ArchitectureConfig = field(default_factory=ArchitectureConfig) + morphology: MorphologyConfig = field(default_factory=MorphologyConfig) + arena: ArenaConfig = field(default_factory=ArenaConfig) + environment: EnvConfig = field(default_factory=EnvConfig) + simulation: SimulationSettings = field(default_factory=SimulationSettings) diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py new file mode 100644 index 0000000..b677b31 --- /dev/null +++ b/src/brittle_star_project/configs/register_configs.py @@ -0,0 +1,40 @@ +from hydra.core.config_store import ConfigStore + +from experiment_logger.config_logger import LoggingConfig +from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_ppo import PPOConfig +from brittle_star_project.configs.config_architecture import ( + CentralizedConfig, + DecentralizedConfig, +) +from brittle_star_project.configs.config_simulation import SimulationSettings +from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig +from brittle_star_project.configs.main_config import BrittleStarConfig + + +def register_configs() -> None: + """Register all dataclasses with Hydra's ConfigStore. + + This must be called before hydra.main() processes the config, ensuring + every structured config is validated against its Python schema. Typos in + YAML keys will raise ConfigAttributeError at startup. + """ + cs = ConfigStore.instance() + + # Root schema + cs.store(name="brittle_star_config", node=BrittleStarConfig) + + # Sub-config groups — each group corresponds to a configs/ subdirectory. + cs.store(group="experiment", name="base_experiment", node=ExperimentConfig) + cs.store(group="logging", name="base_logging", node=LoggingConfig) + cs.store(group="ppo", name="base_ppo", node=PPOConfig) + + # Architecture variants — swap via CLI: architecture=decentralized + cs.store(group="architecture", name="centralized_schema", node=CentralizedConfig) + cs.store(group="architecture", name="decentralized_schema", node=DecentralizedConfig) + + # Environment configs + cs.store(group="morphology", name="base_morphology", node=MorphologyConfig) + cs.store(group="arena", name="base_arena", node=ArenaConfig) + cs.store(group="environment", name="base_environment", node=EnvConfig) + cs.store(group="simulation", name="base_simulation", node=SimulationSettings) diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py deleted file mode 100644 index 036b44f..0000000 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ /dev/null @@ -1,116 +0,0 @@ -from dataclasses import dataclass - -import jax - - -@jax.tree_util.register_dataclass -@dataclass -class PPOArgs: - """ - source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py - """ - - # path to environment config file, if None, use default config - env_config_path: str | None = None - - # path to hyperparameter config file (yaml), if None, use default config - hyperparameter_config_path: str | None = None - - # the name of this experiment - exp_name: str = "brittle_star_ppo" - - # the directory to save the experiment results - run_dir: str | None = None - - # seed of the experiment - seed: int = 1 - - # if toggled, `torch.backends.cudnn.deterministic=False` - torch_deterministic: bool = True - - # if toggled, cuda will be enabled by default - cuda: bool = True - - # if toggled, this experiment will be tracked with Weights and Biases - track: bool = False - - # the wandb's project name - wandb_project_name: str = "PPO-Modularity" - - # the entity (team) of wandb's project - wandb_entity: str | None = "SEL3-2026-Groep-4" - - # whether to capture videos of the agent performances (check out `videos` folder) - capture_video: bool = False - - # 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 - - # the user or org name of the model repository from the Hugging Face Hub - hf_entity: str = "" - - # ==== Algorithm specific dataclasses ==== - - # total timesteps of the experiments - total_timesteps: int = 10000000 - - # the learning rate of the optimizer - learning_rate: float = 2.5e-4 - - # the number of parallel game environments - num_envs: int = 100 - - # the number of steps to run in each environment per policy rollout - num_steps: int = 128 - - # Toggle learning rate annealing for policy and value networks - anneal_lr: bool = True - - # the discount factor gamma - gamma: float = 0.99 - - # the lambda for the general advantage estimation - gae_lambda: float = 0.95 - - # the number of mini-batches - num_minibatches: int = 4 - - # the K epochs to update the policy - update_epochs: int = 4 - - # Toggles advantages normalization - norm_adv: bool = True - - # the surrogate clipping coefficient - clip_coef: float = 0.1 - - # Toggles whether or not to use a clipped loss for the value function, as per the paper. - clip_vloss: bool = True - - # coefficient of the entropy - ent_coef: float = 0.01 - - # coefficient of the value function - vf_coef: float = 0.5 - - # the maximum norm for the gradient clipping - max_grad_norm: float = 0.5 - - # the target KL divergence threshold - target_kl: float | None = None - - # ==== to be filled in runtime ==== - # the batch size (computed in runtime) - batch_size: int = 0 - - # the mini-batch size (computed in runtime) - minibatch_size: int = 0 - - # the number of iterations (computed in runtime) - num_iterations: int = 0 diff --git a/src/brittle_star_project/dataclasses/__init__.py b/src/brittle_star_project/dataclasses/__init__.py index 3501b2b..f2d2ad1 100644 --- a/src/brittle_star_project/dataclasses/__init__.py +++ b/src/brittle_star_project/dataclasses/__init__.py @@ -1,8 +1,6 @@ -from .PPOArgs import PPOArgs from .EpisodeStatistics import EpisodeStatistics __all__ = [ - "PPOArgs", "EpisodeStatistics", ] diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index b143d9c..a5175a7 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -1,14 +1,11 @@ import jax import jax.numpy as jnp -from brittle_star_project import ( - EnvConfig, - BrittleStarEnvFactory, - MorphologyConfig, - ArenaConfig, - Backend, -) -from brittle_star_project.environment import from_file +from experiment_logger import get_logger +from .env_config import EnvConfig, MorphologyConfig, ArenaConfig +from .env_types import Backend +from .factory import BrittleStarEnvFactory +from .padded_obs_wrapper import compute_padding_masks, pad_observations_batched class BrittleStarJaxEnvWrapper: @@ -29,14 +26,15 @@ class BrittleStarJaxEnvWrapper: self._backend, self._morphology, self._arena, self._env_config ) + # Pre-compute masks for observation padding + self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm) + self._vectorized_reset = jax.jit(jax.vmap(self._env.reset)) self._vectorized_step = jax.jit(jax.vmap(self._env.step)) self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample)) 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}" @@ -62,7 +60,12 @@ class BrittleStarJaxEnvWrapper: 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) + state = self._vectorized_reset(rng=env_rngs) + + state = state.replace( + observations=pad_observations_batched(state.observations, self._padding_masks) + ) + return state def sample_actions(self): assert self._action_rng is not None, "Call reset() before sample_actions()" @@ -72,7 +75,12 @@ class BrittleStarJaxEnvWrapper: return self._vectorized_action_sample(rng=jnp.array(sub_rngs)) def step(self, state, action): - return self._vectorized_step(state=state, action=action) + next_state = self._vectorized_step(state=state, action=action) + + next_state = next_state.replace( + observations=pad_observations_batched(next_state.observations, self._padding_masks) + ) + return next_state def close(self): self._env.close() @@ -86,15 +94,6 @@ class BrittleStarJaxEnvWrapper: morphology, arena, env_config, num_envs=num_envs, backend=backend ) - @staticmethod - def from_config( - config_path: str, num_envs: int, backend: Backend = Backend.MJX - ) -> "BrittleStarJaxEnvWrapper": - morphology_cfg, arena_cfg, env_cfg = from_file(config_path) - return BrittleStarJaxEnvWrapper( - morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend - ) - def __str__(self): morphology_str = str(self._morphology) arena_str = str(self._arena) diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index aad8c35..78896cb 100644 --- a/src/brittle_star_project/environment/__init__.py +++ b/src/brittle_star_project/environment/__init__.py @@ -1,4 +1,4 @@ -from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file +from .env_config import ArenaConfig, EnvConfig, MorphologyConfig from .env_types import Backend, Task from .env_wrapper import BrittleStarEnv, StepResult from .factory import BrittleStarEnvFactory @@ -12,5 +12,4 @@ __all__ = [ "BrittleStarEnv", "StepResult", "BrittleStarEnvFactory", - "from_file", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 78083e9..7cb4c21 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -5,24 +5,37 @@ from dataclasses import dataclass, field from .env_types import Task -@dataclass(frozen=True, slots=True) +@dataclass class MorphologyConfig: - num_arms: int = 5 - num_segments_per_arm: int = 4 + """Brittle star morphology configuration. + + segments_per_arm defines the number of segments for each arm. The length of + this list implicitly sets the number of arms. Use 0 segments to represent + a fully amputated arm (e.g., [4, 0, 4, 2, 4] for a 5-arm morphology with + arm 1 removed and arm 3 shortened). + + The upstream biorobot library natively supports per-arm segment counts. + """ + + segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4]) use_p_control: bool = True use_torque_control: bool = False + @property + def num_arms(self) -> int: + return len(self.segments_per_arm) -@dataclass(frozen=True, slots=True) + +@dataclass class ArenaConfig: - size: tuple[float, float] = (10.0, 5.0) + size: list[float] = field(default_factory=lambda: [10.0, 5.0]) sand_ground_color: bool = True attach_target: bool = True wall_height: float = 1.5 wall_thickness: float = 0.1 -@dataclass(frozen=True, slots=True) +@dataclass class EnvConfig: """Shared environment settings. @@ -31,13 +44,13 @@ class EnvConfig: task: Task = Task.DIRECTED_LOCOMOTION - simulation_time: float = 5.0 + simulation_time: float = 10000.0 num_physics_steps_per_control_step: int = 10 time_scale: int = 2 camera_ids: list[int] = field(default_factory=lambda: [0, 1]) # (height, width) - render_size: tuple[int, int] = (480, 640) + render_size: list[int] = field(default_factory=lambda: [480, 640]) joint_randomization_noise_scale: float = 0.0 @@ -47,16 +60,3 @@ class EnvConfig: # Light escape # Per docs in upstream env config: integer factors of 200. light_perlin_noise_scale: int = 0 - - -def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]: - """Load configurations from a YAML file.""" - import yaml - - with open(path, "r") as f: - config_dict = yaml.safe_load(f) - - morphology = MorphologyConfig(**config_dict.get("morphology", {})) - arena = ArenaConfig(**config_dict.get("arena", {})) - env = EnvConfig(**config_dict.get("env", {})) - return morphology, arena, env diff --git a/src/brittle_star_project/environment/factory.py b/src/brittle_star_project/environment/factory.py index 1a891ea..523feb2 100644 --- a/src/brittle_star_project/environment/factory.py +++ b/src/brittle_star_project/environment/factory.py @@ -22,7 +22,7 @@ class BrittleStarEnvFactory: spec = default_brittle_star_morphology_specification( num_arms=config.num_arms, - num_segments_per_arm=config.num_segments_per_arm, + num_segments_per_arm=list(config.segments_per_arm), use_p_control=config.use_p_control, use_torque_control=config.use_torque_control, ) diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py new file mode 100644 index 0000000..3f22038 --- /dev/null +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -0,0 +1,108 @@ +"""Observation padding wrapper for amputated brittle star morphologies. + +When using a centralized controller, the global observation vector must remain +a constant size regardless of how many segments are amputated. This wrapper pads +the observation dictionary values with zeros using spatial insertion so that the +flattened observation maintains the correct physical mapping to the neural network. +""" + +from __future__ import annotations + +from typing import Any +import jax.numpy as jnp + +# Observation keys whose size scales with the number of joints (2 per segment). +_JOINT_SCALED_KEYS = frozenset( + { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + } +) + +# Observation keys whose size scales with the number of segments (1 per segment). +_SEGMENT_SCALED_KEYS = frozenset( + { + "segment_contact", + } +) + + +def compute_padding_masks( + segments_per_arm: tuple[int, ...], + reference_segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4), +) -> dict[str, Any]: + """Pre-compute boolean masks for spatial insertion of observations. + + Args: + segments_per_arm: The current (possibly amputated) morphology. + reference_segments_per_arm: The full morphology that defines the expected size. + + Returns: + A dict containing 1D boolean masks and target sizes. + """ + if len(segments_per_arm) != len(reference_segments_per_arm): + raise ValueError( + f"Morphology mismatch: current has {len(segments_per_arm)} arms, " + f"but reference requires {len(reference_segments_per_arm)} arms." + ) + + mask_1x = [] + mask_2x = [] + + for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)): + if not (0 <= actual <= ref): + raise ValueError( + f"Invalid amputation at arm {arm_idx}: " + f"actual segments ({actual}) must be between 0 and reference ({ref})." + ) + # 1x scaling (e.g., contacts: 1 value per segment) + # 1x scaling (e.g., contacts: 1 value per segment) + mask_1x.extend([True] * actual + [False] * (ref - actual)) + # 2x scaling (e.g., joints: 2 values per segment) + mask_2x.extend([True] * (actual * 2) + [False] * ((ref - actual) * 2)) + + return { + "mask_1x": jnp.array(mask_1x, dtype=bool), + "mask_2x": jnp.array(mask_2x, dtype=bool), + "target_size_1x": sum(reference_segments_per_arm), + "target_size_2x": sum(reference_segments_per_arm) * 2, + } + + +def pad_observation( + obs: dict[str, Any], + masks: dict[str, Any], +) -> dict[str, Any]: + """Pad an observation dict using spatial insertion.""" + padded = {} + for key, value in obs.items(): + if key in _JOINT_SCALED_KEYS: + out = jnp.zeros(masks["target_size_2x"], dtype=value.dtype) + padded[key] = out.at[masks["mask_2x"]].set(value) + elif key in _SEGMENT_SCALED_KEYS: + out = jnp.zeros(masks["target_size_1x"], dtype=value.dtype) + padded[key] = out.at[masks["mask_1x"]].set(value) + else: + padded[key] = value + return padded + + +def pad_observations_batched( + obs: dict[str, Any], + masks: dict[str, Any], +) -> dict[str, Any]: + """Pad a batched observation dict (leading batch dimension) using spatial insertion.""" + padded = {} + for key, value in obs.items(): + batch_size = value.shape[0] + if key in _JOINT_SCALED_KEYS: + out = jnp.zeros((batch_size, masks["target_size_2x"]), dtype=value.dtype) + padded[key] = out.at[:, masks["mask_2x"]].set(value) + elif key in _SEGMENT_SCALED_KEYS: + out = jnp.zeros((batch_size, masks["target_size_1x"]), dtype=value.dtype) + padded[key] = out.at[:, masks["mask_1x"]].set(value) + else: + padded[key] = value + return padded diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index cf6c69e..1ea37a8 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -99,6 +99,7 @@ def get_action_and_value( hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x) hidden_sensor = message_passer(hidden_sensor) mean, log_std = actor_apply(params["actor_params"], hidden_sensor) + log_std = jnp.clip(log_std, -5, 2) std = jnp.exp(log_std) logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) @@ -142,6 +143,7 @@ def ppo_loss( pg_loss1 = -mb_advantages * ratio pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() + v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() entropy_loss = entropy.mean() loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 6bc94ef..e5ecca6 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -13,7 +13,8 @@ from flax.training.train_state import TrainState from experiment_logger import get_logger -from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.dataclasses import EpisodeStatistics from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper from brittle_star_project.MLPs.mlps import ( Actor, @@ -24,6 +25,33 @@ from brittle_star_project.MLPs.mlps import ( ) from brittle_star_project.ppo import PPO +# TODO: move to config +_ALLOWED_OBS_KEYS = { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + "disk_position", + "disk_rotation", + "disk_linear_velocity", + "disk_angular_velocity", + "unit_xy_direction_to_target", + "xy_distance_to_target", +} +# TODO: clip scaled reward? + + +@jax.jit +def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray: + """Extract xy_distance_to_target for all environments.""" + # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...) + return obs_dict["xy_distance_to_target"].squeeze(-1) # shape: (num_envs,) + + +@jax.jit +def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: + return jnp.clip(action, low, high) + def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> float: var_returns = jnp.var(returns) @@ -37,11 +65,25 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear return learning_rate * frac +@jax.jit +def _normalize_obs(obs, mean, var, eps=1e-8): + return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) + + @jax.jit def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: - return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))( - obs_dict - ) + """Convert the raw observation dict → flat array, filtering unwanted keys.""" + + def _filter_and_flatten(o: dict) -> jnp.ndarray: + values = [] + for key in sorted(o.keys()): + if key in _ALLOWED_OBS_KEYS: # TODO: NORMALIZATION or .. of observations?? + v = o[key] + if v.size > 0: + values.append(jnp.asarray(v).flatten()) + return jnp.concatenate(values) + + return jax.vmap(_filter_and_flatten)(obs_dict) def _get_action_and_value_noise( @@ -52,6 +94,8 @@ def _get_action_and_value_noise( agent_state: TrainState, next_obs: jnp.ndarray, key: jax.random.PRNGKey, + action_low, + action_high, ): hidden = sensor.apply(agent_state.params["sensor_params"], next_obs) hidden_critic = feature_extractor.apply( @@ -59,13 +103,16 @@ def _get_action_and_value_noise( ) mean, log_std = actor.apply(agent_state.params["actor_params"], hidden) + log_std = jnp.clip(log_std, -5, 2) key, subkey = jax.random.split(key) noise = jax.random.normal(subkey, shape=mean.shape) std = jnp.exp(log_std) - action = mean + noise * std - logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) + raw_action = mean + noise * std + clipped_action = _clip_action(raw_action, action_low, action_high) + logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) value = critic.apply(agent_state.params["critic_params"], hidden_critic) - return action, logprob, value.squeeze(-1), key + + return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key def _step_once( @@ -76,23 +123,28 @@ def _step_once( feature_extractor: GenericDenseLayersWithActivation, actor: Actor, critic: OneDenseLayerMLP, + action_low, + action_high, ): agent_state, episode_stats, obs, done, key, env_state = carry - action, logprob, value, key = _get_action_and_value_noise( - sensor, feature_extractor, actor, critic, agent_state, obs, key + clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise( + sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high ) episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn( - episode_stats, env_state, action + episode_stats, env_state, clipped_action ) storage = Storage( obs=obs, - actions=action, + actions=raw_action, + raw_actions=raw_action, logprobs=logprob, dones=done, values=value, rewards=reward, + means=mean, + stds=std, returns=jnp.zeros_like(reward), advantages=jnp.zeros_like(reward), ) @@ -103,6 +155,8 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn): next_env_state = env_step_fn(env_state, action) reward = next_env_state.reward + reward *= 20000 + reward = jnp.clip(reward, -10, 10) terminated = next_env_state.terminated truncated = next_env_state.truncated done = terminated | truncated @@ -140,6 +194,8 @@ def _rollout_jit( feature_extractor: GenericDenseLayersWithActivation, actor: Actor, critic: OneDenseLayerMLP, + action_low, + action_high, ): (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan( partial( @@ -149,6 +205,8 @@ def _rollout_jit( actor=actor, critic=critic, env_step_fn=step_env_fn, + action_low=action_low, + action_high=action_high, ), (agent_state, episode_stats, next_obs, next_done, key, env_state), (), @@ -191,7 +249,9 @@ def _compute_gae_jit( (dones[1:], values[1:], values[:-1], storage.rewards), reverse=True, ) - return storage.replace(advantages=advantages, returns=advantages + storage.values) + returns = advantages + storage.values + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) + return storage.replace(advantages=advantages, returns=returns) @dataclass @@ -210,14 +270,23 @@ class TrainingMeasurements: class PPOTrainer: - def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str): - self.args = args + def __init__( + self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str + ): + self.cfg = cfg + self.ppo = cfg.ppo + self.experiment = cfg.experiment + self.logging_cfg = cfg.logging self.env = env self.run_dir = run_dir self.run_name = run_name self.logger = get_logger() - self.key = jax.random.PRNGKey(args.seed) + # Derived runtime fields + self.batch_size = self.ppo.num_envs * self.ppo.num_steps + self.num_iterations = self.ppo.total_timesteps // self.batch_size + + self.key = jax.random.PRNGKey(self.experiment.seed) self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent() self.sensor.apply = jax.jit(self.sensor.apply) @@ -225,29 +294,34 @@ class PPOTrainer: self.actor.apply = jax.jit(self.actor.apply) self.critic.apply = jax.jit(self.critic.apply) + action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) + action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32) + self._rollout_jit = jax.jit( partial( _rollout_jit, - max_steps=self.args.num_steps, + max_steps=self.ppo.num_steps, step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step), sensor=self.sensor, feature_extractor=self.feature_extractor, actor=self.actor, critic=self.critic, + action_low=action_low, + action_high=action_high, ) ) self._compute_gae_jit = jax.jit( partial( _compute_gae_jit, - num_envs=self.args.num_envs, - gamma=self.args.gamma, - gae_lambda=self.args.gae_lambda, + num_envs=self.ppo.num_envs, + gamma=self.ppo.gamma, + gae_lambda=self.ppo.gae_lambda, feature_extractor=self.feature_extractor, critic=self.critic, ) ) - self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor) + self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor) self.agent_state = self._init_agent_state() @@ -256,16 +330,16 @@ class PPOTrainer: self._init_random() def _init_random(self): - self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}") + self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}") - random.seed(self.args.seed) - np.random.seed(self.args.seed) + random.seed(self.experiment.seed) + np.random.seed(self.experiment.seed) def _init_agent(self): self.logger.info("[AGENT]: Initializing agent...") - sensor = GenericDenseLayersWithActivation() - feature_extractor = GenericDenseLayersWithActivation() + sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) + feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) actor = Actor(action_dim=self.env.single_action_space.shape[0]) critic = OneDenseLayerMLP() return sensor, feature_extractor, actor, critic @@ -277,15 +351,11 @@ class PPOTrainer: self.key, 5 ) - sample_obs = jnp.concatenate( - [ - v.flatten() - for v in self.env.single_observation_space.sample( - rng=jax.random.PRNGKey(0) - ).values() - if v.size > 0 - ] - ) + dummy_reset = self.env.reset(seed=0) + sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0] # take first env + self.obs_mean = jnp.zeros((len(sample_obs),)) + self.obs_var = jnp.ones((len(sample_obs),)) + self.obs_count = 1e-4 sensor_params = self.sensor.init(sensor_key, sample_obs) feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs) actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs)) @@ -299,17 +369,17 @@ class PPOTrainer: AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params) ), tx=optax.chain( - optax.clip_by_global_norm(self.args.max_grad_norm), + optax.clip_by_global_norm(self.ppo.max_grad_norm), optax.inject_hyperparams(optax.adam)( learning_rate=partial( _linear_schedule, - minibatch_count=self.args.num_minibatches, - update_epochs=self.args.update_epochs, - num_iterations=self.args.num_iterations, - learning_rate=self.args.learning_rate, + minibatch_count=self.ppo.num_minibatches, + update_epochs=self.ppo.update_epochs, + num_iterations=self.num_iterations, + learning_rate=self.ppo.learning_rate, ) - if self.args.anneal_lr - else self.args.learning_rate, + if self.ppo.anneal_lr + else self.ppo.learning_rate, eps=1e-5, ), ), @@ -319,12 +389,31 @@ class PPOTrainer: self.logger.info("[EPISODE STATS]: Initializing episode stats...") return EpisodeStatistics( - 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), + episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32), + episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), + returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32), + returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), ) + def _update_obs_stats(self, obs: jnp.ndarray): + batch_mean = jnp.mean(obs, axis=0) + batch_var = jnp.var(obs, axis=0) + batch_count = obs.shape[0] + + delta = batch_mean - self.obs_mean + total_count = self.obs_count + batch_count + + new_mean = self.obs_mean + delta * batch_count / total_count + + m_a = self.obs_var * self.obs_count + m_b = batch_var * batch_count + M2 = m_a + m_b + delta**2 * self.obs_count * batch_count / total_count + new_var = M2 / total_count + + self.obs_mean = new_mean + self.obs_var = new_var + self.obs_count = total_count + def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]: return self._rollout_jit( self.agent_state, @@ -350,7 +439,39 @@ class PPOTrainer: start_time, iteration_time_start, training_measurements, + storage, + next_obs, + xy_distance, ): + data = jax.device_get( + { + "rewards": storage.rewards[0], + "values": storage.values[0], + "returns": storage.returns[0], + "advantages": storage.advantages[0], + "actions": storage.actions[0], + "raw_actions": storage.raw_actions[0], + "means": storage.means[0], + "stds": storage.stds[0], + "logprobs": storage.logprobs[0], + } + ) + + storage_metrics = { + "rollout/env0/return_mean": float(np.mean(data["returns"])), + "rollout/env0/advantage_mean": float(np.mean(data["advantages"])), + "rollout/env0/value_mean": float(np.mean(data["values"])), + "rollout/env0/value_vs_return_diff": float(np.mean(data["values"] - data["returns"])), + "rollout/env0/reward_mean": float(np.mean(data["rewards"])), + "rollout/env0/mean_mean": float(np.mean(data["means"])), + "rollout/env0/logprob_mean": float(np.mean(data["logprobs"])), + "rollout/env0/action_mean": float(np.mean(data["actions"])), + "rollout/env0/raw_action_mean": float(np.mean(data["raw_actions"])), + } + + for i in range(len(xy_distance)): + storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i]) + metrics = { "charts/avg_episodic_return": training_measurements.avg_episodic_return, "charts/avg_episodic_length": np.mean( @@ -371,8 +492,9 @@ class PPOTrainer: "losses/loss": training_measurements.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.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start) ), + **storage_metrics, } self.logger.log(metrics, step=global_step) @@ -443,6 +565,7 @@ class PPOTrainer: avg_terminated_length=avg_terminated_length, avg_truncated_length=avg_truncated_length, ), + storage, ) def _close(self): @@ -451,8 +574,14 @@ class PPOTrainer: def _save_model(self, model_path: str): self.logger.info("[SAVE]: Saving the final model...") + from dataclasses import asdict as _asdict + + config_dict = { + "experiment": _asdict(self.experiment), + "ppo": _asdict(self.ppo), + } params = [ - vars(self.args), + config_dict, [ self.agent_state.params["sensor_params"], self.agent_state.params["actor_params"], @@ -464,8 +593,7 @@ class PPOTrainer: def train(self): """ - Train the PPO agent for a specified number of iterations - (passed through PPOArgs in constructor). + Train the PPO agent for a specified number of iterations. Closes the environment at the end of training. """ self.logger.info(f"running name: {self.run_name}") @@ -473,47 +601,58 @@ class PPOTrainer: 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) + env_state = self.env.reset(seed=self.experiment.seed) next_obs = _convert_obs_dict_to_array(env_state.observations) - next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_) + next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_) self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}") global_step = 0 start_time = time.time() - iter_bar = self.logger.progress_bar(range(1, self.args.num_iterations + 1)) + iter_bar = self.logger.progress_bar(range(1, self.num_iterations + 1)) for iteration in iter_bar: iteration_time_start = time.time() - env_state, next_obs, next_done, training_measurements = self._step( + env_state, next_obs, next_done, training_measurements, storage = self._step( env_state, next_obs, next_done, iteration=iteration ) + self._update_obs_stats(next_obs) + next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var) - global_step += self.args.num_steps * self.args.num_envs + xy_distance = _get_xy_distance_to_target(env_state.observations) + + global_step += self.ppo.num_steps * self.ppo.num_envs self._log( global_step, self.episode_stats, start_time, iteration_time_start, training_measurements, + storage, + next_obs, + xy_distance, ) sps = int(global_step / (time.time() - start_time)) - remaining_steps = self.args.total_timesteps - global_step + remaining_steps = self.ppo.total_timesteps - global_step eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 eta_str = str(datetime.timedelta(seconds=eta_seconds)) self.logger.log_non_interactive( - f"Iteration {iteration}/{self.args.num_iterations} | " - f"Step {global_step}/{self.args.total_timesteps} | " + f"Iteration {iteration}/{self.num_iterations} | " + f"Step {global_step}/{self.ppo.total_timesteps} | " f"SPS {sps} | " f"Return {training_measurements.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" + if getattr(self.cfg.experiment, "debug_sanity", False): + self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch") + break + + if self.logging_cfg.save_model: + model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model" self._save_model(model_path=model_path) self._close() diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index 64d4be4..53e57c0 100644 --- a/src/experiment_logger/__init__.py +++ b/src/experiment_logger/__init__.py @@ -4,8 +4,8 @@ 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.config_utils import load_yaml_config +from experiment_logger.unified_logger import UnifiedLogger, get_logger, init_logger from experiment_logger.simple_logger import SimpleLogger from experiment_logger.wandb_utils import finish_wandb, init_wandb @@ -13,9 +13,9 @@ __all__ = [ "UnifiedLogger", "SimpleLogger", "get_logger", + "init_logger", "init_wandb", "finish_wandb", "load_yaml_config", - "merge_config_with_cli", ] __version__ = "0.1.0" diff --git a/src/experiment_logger/config_logger.py b/src/experiment_logger/config_logger.py new file mode 100644 index 0000000..5d39b8e --- /dev/null +++ b/src/experiment_logger/config_logger.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class LoggingConfig: + track: bool = False + wandb_project_name: str = "PPO-Modularity" + wandb_entity: Optional[str] = "SEL3-2026-Groep-4" + capture_video: bool = False + save_model: bool = True + checkpoint_frequency: int = 100 + upload_model: bool = False + hf_entity: str = "" diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index 4c5b79f..80d1748 100644 --- a/src/experiment_logger/config_utils.py +++ b/src/experiment_logger/config_utils.py @@ -1,15 +1,12 @@ """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") @@ -24,7 +21,7 @@ def load_yaml_config(config_path: str) -> Dict[str, Any]: if config is None: return {} - log.info(f"Loaded configuration from: {config_path}") + get_logger().info(f"Loaded configuration from: {config_path}") return config @@ -35,7 +32,7 @@ def save_yaml_config(config: Dict[str, Any], config_path: str): 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}") + get_logger().info(f"Saved configuration to: {config_path}") def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: @@ -66,82 +63,21 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: 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}") + get_logger().warning(f"Could not convert {key}={value} to {field.type}: {e}") filtered_config[key] = value else: - log.warning(f"Unknown configuration parameter: {key}") + get_logger().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}:") + get_logger().info(f"{title}:") if is_dataclass(config): for field in fields(config): value = getattr(config, field.name) - log.info(f" {field.name}: {value}") + get_logger().info(f" {field.name}: {value}") else: for key, value in vars(config).items(): - log.info(f" {key}: {value}") + get_logger().info(f" {key}: {value}") diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 5cfb3f1..2f5d06f 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -6,9 +6,7 @@ This logger ensures all experimental data is preserved by writing to: 3. stdout (for real-time monitoring) """ -import datetime import logging -import subprocess import yaml import sys import time @@ -21,38 +19,67 @@ import numpy as np from experiment_logger.wandb_utils import finish_wandb, init_wandb -# Global singleton storage -_global_logger = None +# Global storage for the active logger and the proxy singleton +_active_logger: Optional[Any] = None +_proxy_instance: Optional["LoggerProxy"] = 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" +def get_logger() -> "LoggerProxy": + """Retrieve the global LoggerProxy. - timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - generic_name = f"{timestamp_str}_{commit_hash}_brittle_star" + This should be used for all logging calls. It returns a proxy that + delegates to the active logger (defaulting to a SimpleLogger until + init_logger is called). + """ + global _proxy_instance, _active_logger + if _proxy_instance is None: + if _active_logger is None: + # Fallback to SimpleLogger to avoid premature directory creation + from experiment_logger.simple_logger import SimpleLogger - # 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}") + _active_logger = SimpleLogger(run_name="pre_init") - return _global_logger + _proxy_instance = LoggerProxy() + + return _proxy_instance + + +def init_logger(**kwargs) -> "UnifiedLogger": + """Initialize the full UnifiedLogger and set it as the active logger. + + This should be called once the configuration is ready. It will create + the output directories and set up all logging backends. + """ + global _active_logger + logger = UnifiedLogger(**kwargs) + _active_logger = logger + return logger + + +class LoggerProxy: + """Proxy that delegates all method calls to the active logger instance. + + This allows the logger to be swapped out (e.g., from a SimpleLogger to + a UnifiedLogger) without any clients needing to update their references. + """ + + def _get_logger(self) -> Any: + global _active_logger + if _active_logger is None: + # This shouldn't normally happen since get_logger handles it + from experiment_logger.simple_logger import SimpleLogger + + _active_logger = SimpleLogger(run_name="pre_init_fallback") + return _active_logger + + def __getattr__(self, name: str) -> Any: + return getattr(self._get_logger(), name) + + def __enter__(self): + return self._get_logger().__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + return self._get_logger().__exit__(exc_type, exc_val, exc_tb) class UnifiedLogger: @@ -68,7 +95,6 @@ class UnifiedLogger: use_wandb: bool = True, save_code: bool = True, log_level: int = logging.INFO, - _set_as_global: bool = True, ): """Initialize the unified logger. @@ -80,7 +106,6 @@ class UnifiedLogger: base_dir: Base directory for local storage use_wandb: Whether to use WandB logging save_code: Whether to save code to WandB - _set_as_global: Internal flag to override the global singleton """ self.run_name = run_name self.config = config @@ -119,11 +144,6 @@ class UnifiedLogger: 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() diff --git a/tests/.gitkeep b/tests/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/tests/test_config.py b/tests/test_config.py deleted file mode 100644 index 212e4b4..0000000 --- a/tests/test_config.py +++ /dev/null @@ -1,38 +0,0 @@ -"""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/tests/test_configs.py b/tests/test_configs.py new file mode 100644 index 0000000..0a6a651 --- /dev/null +++ b/tests/test_configs.py @@ -0,0 +1,48 @@ +from pathlib import Path +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs + +# Registration must happen before composition to enable validation against schemas +register_configs() + + +def test_config_composition_centralized(): + """Test that the centralized configuration composes and validates correctly.""" + config_dir = str(Path(__file__).parent.parent / "configs") + with initialize_config_dir(version_base="1.3", config_dir=config_dir): + # We compose the config; it follows main_config.yaml + cfg = compose(config_name="main_config", overrides=["architecture=centralized"]) + + # Merge with the structured schema and convert to a real dataclass instance + structured_cfg = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg) + ) + + # Basic assertions + assert structured_cfg.architecture.name == "centralized" + assert structured_cfg.architecture.propagator is None + assert isinstance(structured_cfg.ppo.learning_rate, float) + assert structured_cfg.ppo.learning_rate > 0 + + +def test_config_composition_decentralized(): + """Test that the decentralized configuration composes and validates correctly.""" + config_dir = str(Path(__file__).parent.parent / "configs") + with initialize_config_dir(version_base="1.3", config_dir=config_dir): + cfg = compose(config_name="main_config", overrides=["architecture=decentralized"]) + + # Merge and convert to dataclass instance + structured_cfg = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg) + ) + + # Basic assertions + assert structured_cfg.architecture.name == "decentralized" + assert isinstance(structured_cfg.ppo.learning_rate, float) + assert structured_cfg.ppo.learning_rate > 0 + + # Decentralized specifics + assert hasattr(structured_cfg.architecture, "message_passing_steps") + assert structured_cfg.architecture.message_passing_steps > 0 diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py new file mode 100644 index 0000000..6b435cf --- /dev/null +++ b/tests/test_morphology_render.py @@ -0,0 +1,65 @@ +import pytest +import os +import sys + +# CRITICAL for headless cross-platform testing (devcontainers etc) +if sys.platform == "linux" and "DISPLAY" not in os.environ and "WAYLAND_DISPLAY" not in os.environ: + os.environ.setdefault("MUJOCO_GL", "egl") + +import mujoco +from PIL import Image +from brittle_star_project.environment.env_config import EnvConfig, MorphologyConfig, ArenaConfig +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper + + +@pytest.mark.skipif(os.getenv("CI") == "true", reason="No OpenGL display in CI") +def test_render_morphologies(): + base_dir = "runs/renders" + os.makedirs(base_dir, exist_ok=True) + + # --- 1. Full 5-Arm Morphology --- + morph_full = MorphologyConfig(segments_per_arm=[4, 4, 4, 4, 4]) + env_full = BrittleStarJaxEnvWrapper( + morphology=morph_full, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1 + ) + state_full = env_full.reset(seed=0) + + model_full = state_full.mj_model + data_full = state_full.mj_data + + # 1. Compute forward kinematics so geoms are correctly positioned + mujoco.mj_forward(model_full, data_full) + + # 2. Render using the environment's primary camera (camera=0) + renderer_full = mujoco.Renderer(model=model_full) + renderer_full.update_scene(data_full, camera=1) + pixels_full = renderer_full.render() + image_path = os.path.join(base_dir, "full_5_arm.png") + Image.fromarray(pixels_full).save(image_path) + print(f"Generated full morphology render: {image_path}") + + # --- 2. Partially Amputated Morphology --- + morph_amp = MorphologyConfig(segments_per_arm=[4, 0, 4, 2, 4]) + env_amp = BrittleStarJaxEnvWrapper( + morphology=morph_amp, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1 + ) + state_amp = env_amp.reset(seed=0) + + model_amp = state_amp.mj_model + data_amp = state_amp.mj_data + + # Compute forward kinematics + mujoco.mj_forward(model_amp, data_amp) + + renderer_amp = mujoco.Renderer(model=model_amp) + renderer_amp.update_scene(data_amp, camera=1) + pixels_amp = renderer_amp.render() + image_path = os.path.join(base_dir, "amputated_arm.png") + Image.fromarray(pixels_amp).save(image_path) + print(f"Generated amputated morphology render: {image_path}") + + print("Morphology render test successful!") + + +if __name__ == "__main__": + test_render_morphologies() diff --git a/tests/test_network_shapes.py b/tests/test_network_shapes.py new file mode 100644 index 0000000..32e9898 --- /dev/null +++ b/tests/test_network_shapes.py @@ -0,0 +1,67 @@ +import jax +import jax.numpy as jnp +from brittle_star_project.environment.padded_obs_wrapper import ( + compute_padding_masks, + pad_observations_batched, +) + +# We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py +from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP + + +def test_centralized_forward_pass_with_padding(): + batch_size = 2 + + # 1. Simulate Amputated Observation [4, 0, 4, 2, 4] -> 14 segments total + # 14 segments * 2 = 28 joints + amputated_obs = { + "joint_position": jnp.zeros((batch_size, 28)), + "joint_velocity": jnp.zeros((batch_size, 28)), + "segment_contact": jnp.zeros((batch_size, 14)), + } + + # 2. Pad Observation using the boolean scattering wrapper + masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4)) + padded_obs = pad_observations_batched(amputated_obs, masks) + + # Assertions to ensure padding sizes are correct (40 joints, 20 segments) + assert padded_obs["joint_position"].shape == (batch_size, 40), "Padding failed for joint keys" + assert padded_obs["segment_contact"].shape == (batch_size, 20), ( + "Padding failed for segment keys" + ) + + # 3. Concatenate for Centralized MLP (simulating the global state vector) + global_state = jnp.concatenate( + [padded_obs["joint_position"], padded_obs["joint_velocity"], padded_obs["segment_contact"]], + axis=-1, + ) + + # 40 + 40 + 20 = 100 dimensions + assert global_state.shape == (batch_size, 100), ( + f"Expected global state shape (2, 100), got {global_state.shape}" + ) + + # 4. Initialize dummy networks (40 actuators for the max morphology output) + actor = Actor(action_dim=40) + critic = OneDenseLayerMLP() # Acts as the centralized critic + + rng = jax.random.PRNGKey(0) + rng_a, rng_c = jax.random.split(rng) + + # Initialize Flax variables + actor_params = actor.init(rng_a, global_state) + critic_params = critic.init(rng_c, global_state) + + # 5. Forward Pass Assertions + action_mean, action_log_std = actor.apply(actor_params, global_state) + value = critic.apply(critic_params, global_state) + + assert action_mean.shape == (batch_size, 40), f"Actor mean shape mismatch: {action_mean.shape}" + assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}" + assert value.shape == (batch_size, 1) or value.shape == (batch_size,), ( + f"Critic value shape mismatch: {value.shape}" + ) + + +if __name__ == "__main__": + test_centralized_forward_pass_with_padding() diff --git a/uv.lock b/uv.lock index 01caa5d..bca7772 100644 --- a/uv.lock +++ b/uv.lock @@ -18,6 +18,7 @@ dependencies = [ { name = "evosax" }, { name = "flax" }, { name = "gymnasium" }, + { name = "hydra-core" }, { name = "ipykernel" }, { name = "jax" }, { name = "matplotlib" }, @@ -30,7 +31,6 @@ dependencies = [ { name = "pyopengl-accelerate" }, { name = "pyyaml" }, { name = "torch" }, - { name = "tyro" }, { name = "wandb" }, { name = "warp-lang" }, ] @@ -57,6 +57,7 @@ requires-dist = [ { name = "evosax", specifier = "==0.2.0" }, { name = "flax", specifier = ">=0.12.2" }, { name = "gymnasium", specifier = ">=1.2.3" }, + { name = "hydra-core", specifier = ">=1.3.2" }, { name = "ipykernel", specifier = "==7.2.0" }, { name = "jax", specifier = "==0.9.0.1" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" }, @@ -71,7 +72,6 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0" }, { name = "tensorboard", marker = "extra == 'analysis'" }, { name = "torch", specifier = ">=2.4.0" }, - { name = "tyro", specifier = ">=1.0.10" }, { name = "wandb", specifier = "==0.24.2" }, { name = "warp-lang" }, ] @@ -111,6 +111,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + [[package]] name = "appnope" version = "0.1.4" @@ -492,15 +498,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ea/97/4f78412f73a9350bc8f934441bae5b68b102c8f4240a7f06b4114b51d6de/dm_tree-0.1.9-cp312-cp312-win_amd64.whl", hash = "sha256:9020a5ce256fcc83aa4bc190cc96dd66e87685db0a6e501b0c06aa492c2e38fc", size = 102022, upload-time = "2025-01-30T20:45:28.701Z" }, ] -[[package]] -name = "docstring-parser" -version = "0.17.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b2/9d/c3b43da9515bd270df0f80548d9944e389870713cc1fe2b8fb35fe2bcefd/docstring_parser-0.17.0.tar.gz", hash = "sha256:583de4a309722b3315439bb31d64ba3eebada841f2e2cee23b99df001434c912", size = 27442, upload-time = "2025-07-21T07:35:01.868Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/55/e2/2537ebcff11c1ee1ff17d8d0b6f4db75873e3b0fb32c2d4a2ee31ecb310a/docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708", size = 36896, upload-time = "2025-07-21T07:35:00.684Z" }, -] - [[package]] name = "dotmap" version = "1.3.30" @@ -753,6 +750,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" }, ] +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + [[package]] name = "identify" version = "2.6.18" @@ -1531,6 +1542,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/11/3f1ee9dce24b41812dd572a037c4436d4d21f759fbe373cc271b0ce98805/nvidia_nvvm-13.2.51-py3-none-win_amd64.whl", hash = "sha256:a4809baaa5429eabe1878853761ce31f0ba15216e2348710b7898dc591f5fc14", size = 56751075, upload-time = "2026-03-09T10:11:09.994Z" }, ] +[[package]] +name = "omegaconf" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, +] + [[package]] name = "opencv-python" version = "4.13.0.92" @@ -2315,18 +2339,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, ] -[[package]] -name = "typeguard" -version = "4.5.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2b/e8/66e25efcc18542d58706ce4e50415710593721aae26e794ab1dec34fb66f/typeguard-4.5.1.tar.gz", hash = "sha256:f6f8ecbbc819c9bc749983cc67c02391e16a9b43b8b27f15dc70ed7c4a007274", size = 80121, upload-time = "2026-02-19T16:09:03.392Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" }, -] - [[package]] name = "typing-extensions" version = "4.15.0" @@ -2361,20 +2373,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] -[[package]] -name = "tyro" -version = "1.0.10" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "docstring-parser" }, - { name = "typeguard" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/30/c1/0a5850badd3f18373d6a0366091638674cec6780b558c1c5b846adea938b/tyro-1.0.10.tar.gz", hash = "sha256:2822eacac963a4922bf7eafe3b156a1f0f7fe8e34148202987581224f25565c2", size = 481084, upload-time = "2026-03-18T08:24:17.307Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/05/be/a0b4c9fa64999a2e337cbefcdedd2e101e8dd88a84e4fa497bd0e4531dc1/tyro-1.0.10-py3-none-any.whl", hash = "sha256:8de87a3a40c8a91f10831f8f0638cd0eed00f0e4de9cd3d561e967f407477210", size = 183433, upload-time = "2026-03-18T08:24:16.012Z" }, -] - [[package]] name = "tzdata" version = "2025.3"