Merge pull request #37 from SELab-3-2026/feat/modular-configurations

.gitignore

@@ -2,6 +2,9 @@
 artifacts/*
 runs/*
 wandb/
+outputs/
+multirun/
+metrics/
 
 # Python-generated files
 __pycache__/

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.
+### Dry-Run Validation
-- `dev_test.yaml`: Fast iteration for development.
+Check if your configuration is valid without starting the simulation:
-- `production_training.yaml`: Full-scale training.
+```bash
-- `personal_template.yaml`: Template for team members to customize.
+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.

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"

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"

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

configs/default_ppo.yaml

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

configs/dev_test.yaml

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

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

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

configs/example.yaml

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

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

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

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

configs/hpc/debug.yaml

@@ -1,20 +0,0 @@
-# Configuration for debug session
-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

configs/hpc/smoke_test.yaml

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

configs/hpc/wandb_expand.yaml

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

configs/hpc/wandb_test.yaml

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

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

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

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}

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

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

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

configs/personal_template.yaml

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

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

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

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

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

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

configs/production_training.yaml

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

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"

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

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

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/debug.yaml \
+    hydra.run.dir="$SCRATCH_RUNDIR" \
-    --hyperparameter-config-path configs/hpc/debug.yaml \
+    ppo=stable \
-    --run-dir "$SCRATCH_RUNDIR"
+    logging=wandb_enabled
 
 echo ">>> Staging out results to $DATA_RUNDIR..."
 cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"

scripts/simulate.py

@@ -1,95 +1,83 @@
+"""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 hydra
+from omegaconf import DictConfig, OmegaConf
 from pathlib import Path
 
 from brittle_star_project import (
-    Backend,
     BrittleStarEnv,
     BrittleStarEnvFactory,
     SimulationConfig,
     simulate_policy,
 )
-from brittle_star_project.environment import from_file
+from brittle_star_project.configs.main_config import BrittleStarConfig
-from brittle_star_project.rl import RLModel  # imports concrete models via rl.__init__
+from brittle_star_project.configs.register_configs import register_configs
+from brittle_star_project.rl import RLModel
 from brittle_star_project.rl.base import get_rl_model_registry
 
 MODEL_BY_NAME = get_rl_model_registry()
 MODEL_OPTIONS = sorted(MODEL_BY_NAME)
 
 
-def parse_args() -> argparse.Namespace:
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
-    p = argparse.ArgumentParser(description="Simulate a trained policy in the MuJoCo viewer.")
+def main(dict_cfg: DictConfig) -> None:
-    p.add_argument(
+    # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly.
-        "--model",
+    config: BrittleStarConfig = OmegaConf.to_object(
-        type=str,
+        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
-        default=None,
-        help="Path to a saved model artifact. If omitted, a model is created from --model-type.",
     )
-    p.add_argument(
-        "--model-type",
-        choices=MODEL_OPTIONS,
-        default="random",
-        help="Which model class to instantiate when --model is omitted.",
-    )
-    p.add_argument(
-        "--backend",
-        choices=[b for b in Backend],
-        default=Backend.MJX,
-    )
-    p.add_argument("--seed", type=int, default=None)
-    return p.parse_args()
 
+    # Use the configurable settings from the simulation group
+    backend = config.simulation.backend
+    model_type = config.simulation.model_type
 
-def main() -> None:
+    # Hydra chdir changes CWD; we map CLI relative paths relative to invocation originally.
-    args = parse_args()
+    model_path = config.simulation.model_path
+    if model_path is not None:
+        model_path = hydra.utils.to_absolute_path(model_path)
 
-    morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml")
+    seed = config.experiment.seed
 
     # ======= ENVIRONMENT SETUP =======
-
-    backend = args.backend
-
     factory = BrittleStarEnvFactory()
-    raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
+    raw_env = factory.create_environment(
-    env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg)
+        backend, config.morphology, config.arena, config.environment
+    )
+    env = BrittleStarEnv(raw_env, backend=backend, config=config.environment)
 
-    seed_for_env = int(args.seed) if args.seed is not None else 0
+    state = env.reset(seed=seed)
-    state = env.reset(seed=seed_for_env)
 
     # ======= MODEL SETUP =======
-
-    # 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)
 
-    if args.model is not None:
+    if model_path is not None:
-        model_path = Path(args.model)
+        policy = RLModel.load(Path(model_path))
-        policy = RLModel.load(model_path)
         if hasattr(policy, "nu"):
             policy.nu = nu
     else:
-        model_cls = MODEL_BY_NAME[str(args.model_type)]
+        model_cls = MODEL_BY_NAME[model_type]
-        policy = model_cls(seed=seed_for_env)
+        policy = model_cls(seed=seed)
         if hasattr(policy, "nu"):
             policy.nu = nu
 
-    # If the policy/model has a `seed` attribute, use the provided seed (or default) to reset it.
+    default_seed = int(getattr(policy, "seed", seed))
-    default_seed = int(getattr(policy, "seed", seed_for_env))
-    if args.seed is not None and hasattr(policy, "reset"):
-        policy.reset(int(args.seed))
 
     # ======= SIMULATION =======
-
     rollout_cfg = SimulationConfig(
         realtime=True,
-        seed=int(args.seed) if args.seed is not None else default_seed,
+        seed=default_seed,
     )
 
     simulate_policy(policy, rollout_cfg, state)
-
     env.close()
 
 
 if __name__ == "__main__":
+    register_configs()
     main()

scripts/train.py

@@ -1,76 +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 init_logger, get_logger
-from experiment_logger import UnifiedLogger
-from experiment_logger.config_utils import merge_config_with_cli, print_config
 
 
-def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper:
+def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper:
-    if config_path is None:
+    """Create the environment using the structured configuration."""
-        return BrittleStarJaxEnvWrapper.default(num_envs=num_envs)
+    return BrittleStarJaxEnvWrapper(
-    return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs)
+        morphology=cfg.morphology,
+        arena=cfg.arena,
+        env_config=cfg.environment,
+        num_envs=cfg.ppo.num_envs,
+    )
 
 
-def parse_args() -> PPOArgs:
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
-    import argparse
+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
+    # 2. Setup run metadata
-    parser = argparse.ArgumentParser(add_help=False)
+    # Hydra changes CWD to the output directory by default.
-    parser.add_argument("--hyperparameter-config-path", type=str, default=None)
+    run_dir = os.getcwd()
-    known_args, _ = parser.parse_known_args()
+    run_name = os.path.basename(run_dir)
 
-    args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path)
+    # 3. Initialize Logger
-    return args
+    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:
+    # 5. Train - pass structured config directly
-    try:
+    ppo_trainer = PPOTrainer(config, env, run_dir, run_name)
-        return (
+    ppo_trainer.train()
-            subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
-        )
-    except (subprocess.CalledProcessError, UnicodeDecodeError):
-        return "none"
 
 
 if __name__ == "__main__":
-    args = parse_args()
+    register_configs()
-
+    main()
-    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)
-    raw_env = env.raw
-
-    torch.backends.cudnn.deterministic = args.torch_deterministic
-
-    ppo_trainer = PPOTrainer(args, env, run_dir, run_name)
-    ppo_trainer.train()

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"

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"

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

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

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)

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)

src/brittle_star_project/dataclasses/PPOArgs.py

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

src/brittle_star_project/dataclasses/__init__.py

@@ -1,8 +1,6 @@
-from .PPOArgs import PPOArgs
 from .EpisodeStatistics import EpisodeStatistics
 
 
 __all__ = [
-    "PPOArgs",
     "EpisodeStatistics",
 ]

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -1,14 +1,11 @@
 import jax
 import jax.numpy as jnp
 
-from brittle_star_project import (
+from experiment_logger import get_logger
-    EnvConfig,
+from .env_config import EnvConfig, MorphologyConfig, ArenaConfig
-    BrittleStarEnvFactory,
+from .env_types import Backend
-    MorphologyConfig,
+from .factory import BrittleStarEnvFactory
-    ArenaConfig,
+from .padded_obs_wrapper import compute_padding_masks, pad_observations_batched
-    Backend,
-)
-from brittle_star_project.environment import from_file
 
 
 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)

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

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
+    """Brittle star morphology configuration.
-    num_segments_per_arm: int = 4
+
+    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.
 
@@ -37,7 +50,7 @@ class EnvConfig:
 
     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

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

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

src/brittle_star_project/trainers/PPOTrainer.py

@@ -11,7 +11,10 @@ import numpy as np
 import optax
 from flax.training.train_state import TrainState
 
-from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
+from experiment_logger import get_logger
+
+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,
@@ -21,7 +24,6 @@ from brittle_star_project.MLPs.mlps import (
     Storage,
 )
 from brittle_star_project.ppo import PPO
-from experiment_logger import get_logger
 
 # TODO: move to config
 _ALLOWED_OBS_KEYS = {
@@ -268,14 +270,23 @@ class TrainingMeasurements:
 
 
 class PPOTrainer:
-    def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str):
+    def __init__(
-        self.args = args
+        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)
@@ -289,7 +300,7 @@ class PPOTrainer:
         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,
@@ -302,15 +313,15 @@ class PPOTrainer:
         self._compute_gae_jit = jax.jit(
             partial(
                 _compute_gae_jit,
-                num_envs=self.args.num_envs,
+                num_envs=self.ppo.num_envs,
-                gamma=self.args.gamma,
+                gamma=self.ppo.gamma,
-                gae_lambda=self.args.gae_lambda,
+                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()
 
@@ -319,10 +330,10 @@ 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)
+        random.seed(self.experiment.seed)
-        np.random.seed(self.args.seed)
+        np.random.seed(self.experiment.seed)
 
     def _init_agent(self):
         self.logger.info("[AGENT]: Initializing agent...")
@@ -358,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,
+                        minibatch_count=self.ppo.num_minibatches,
-                        update_epochs=self.args.update_epochs,
+                        update_epochs=self.ppo.update_epochs,
-                        num_iterations=self.args.num_iterations,
+                        num_iterations=self.num_iterations,
-                        learning_rate=self.args.learning_rate,
+                        learning_rate=self.ppo.learning_rate,
                     )
-                    if self.args.anneal_lr
+                    if self.ppo.anneal_lr
-                    else self.args.learning_rate,
+                    else self.ppo.learning_rate,
                     eps=1e-5,
                 ),
             ),
@@ -378,10 +389,10 @@ class PPOTrainer:
         self.logger.info("[EPISODE STATS]: Initializing episode stats...")
 
         return EpisodeStatistics(
-            episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
+            episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32),
-            episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+            episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
-            returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
+            returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32),
-            returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+            returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
         )
 
     def _update_obs_stats(self, obs: jnp.ndarray):
@@ -481,7 +492,7 @@ 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,
         }
@@ -563,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"],
@@ -576,8 +593,7 @@ class PPOTrainer:
 
     def train(self):
         """
-        Train the PPO agent for a specified number of iterations
+        Train the PPO agent for a specified number of iterations.
-        (passed through PPOArgs in constructor).
         Closes the environment at the end of training.
         """
         self.logger.info(f"running name: {self.run_name}")
@@ -585,16 +601,16 @@ 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()
 
@@ -606,7 +622,7 @@ class PPOTrainer:
 
             xy_distance = _get_xy_distance_to_target(env_state.observations)
 
-            global_step += self.args.num_steps * self.args.num_envs
+            global_step += self.ppo.num_steps * self.ppo.num_envs
             self._log(
                 global_step,
                 self.episode_stats,
@@ -619,20 +635,24 @@ class PPOTrainer:
             )
 
             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"Iteration {iteration}/{self.num_iterations} | "
-                f"Step {global_step}/{self.args.total_timesteps} | "
+                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:
+            if getattr(self.cfg.experiment, "debug_sanity", False):
-            model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model"
+                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()

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.config_utils import load_yaml_config
-from experiment_logger.unified_logger import UnifiedLogger, get_logger
+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"

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

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}")

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 storage for the active logger and the proxy singleton
-_global_logger = None
+_active_logger: Optional[Any] = None
+_proxy_instance: Optional["LoggerProxy"] = None
 
 
-def get_logger() -> "UnifiedLogger":
+def get_logger() -> "LoggerProxy":
-    """Retrieve the global UnifiedLogger. If not initialized, fallback to auto-initialization."""
+    """Retrieve the global LoggerProxy.
-    global _global_logger
-    if _global_logger is None:
-        try:
-            commit_hash = (
-                subprocess.check_output(
-                    ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.STDOUT
-                )
-                .decode("utf-8")
-                .strip()
-            )
-        except Exception:
-            commit_hash = "unknown"
 
-        timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    This should be used for all logging calls. It returns a proxy that
-        generic_name = f"{timestamp_str}_{commit_hash}_brittle_star"
+    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
+            _active_logger = SimpleLogger(run_name="pre_init")
-        _global_logger = UnifiedLogger(
-            run_name=generic_name,
-            config={"auto_initialized": True},
-            use_wandb=False,
-            _set_as_global=False,  # Prevent recursive call inside __init__
-        )
-        _global_logger.warning(f"UnifiedLogger auto-initialized with name: {generic_name}")
 
-    return _global_logger
+        _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()
 

tests/test_config.py

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

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

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

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

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"